Add `eqV2_86M_omat_mp_salex` model (#14)

* refactor external calculators

* add `eqV2_86M_omat_mp_salex`

* change eos test MLIPEnum

* fix bugs for MLIPEnum refactoring

* login huggingface for download ckpt during test

* remove add git credential

* use hf read only token

* rollback `fairchem==1.1.0`

* gracefully xfail for some farichem models

* reverse test installation order

* try adding opt-einsum constraint

* try installing from freeze

* try installing from freeze

* loosen cachepath

* add test pip cache

* bump mace to 0.3.6

* enforce `e3nn==0.4.4`

* add prefect test harness

* change prefect text fixture

* add diatomic curves

* add matgl json

* increase concurrent number

* pytest only on py3.11

* fix system info tuple bug

* use context manager

.github/workflows/test.yaml

@@ -7,17 +7,20 @@ jobs:
     runs-on: ubuntu-latest
 
     strategy:
+      # max-parallel: 2
       matrix:
         python-version: ["3.10", "3.11"]
+        
 
     steps:
     - name: Checkout repository
-      uses: actions/checkout@v2
+      uses: actions/checkout@v4
 
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
+        cache: 'pip'
 
     - name: Install dependencies
       run: |
@@ -25,10 +28,19 @@ jobs:
         pip install torch==2.2.0
         bash scripts/install-pyg.sh
         bash scripts/install-dgl.sh
-        pip install .[test]
         pip install .[mace]
+        pip install .[test]
         pip install "pynanoflann@git+https://github.com/dwastberg/pynanoflann#egg=af434039ae14bedcbb838a7808924d6689274168"
 
+    - name: List dependencies
+      run: pip list
+      
+    - name: Login huggingface
+      env:
+        HF_TOKEN: ${{ secrets.HF_TOKEN_READ_ONLY }}
+      run:
+        huggingface-cli login --token $HF_TOKEN
+
     - name: Run tests
       env:
         PREFECT_API_KEY: ${{ secrets.PREFECT_API_KEY }}

mlip_arena/models/__init__.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import importlib
+from enum import Enum
 from pathlib import Path
 
 import torch
@@ -14,6 +16,17 @@ from torch import nn
 with open(Path(__file__).parent / "registry.yaml", encoding="utf-8") as f:
     REGISTRY = yaml.safe_load(f)
 
+MLIPMap = {}
+
+for model, metadata in REGISTRY.items():
+    try:
+        module = importlib.import_module(f"{__package__}.{metadata['module']}.{metadata['family']}")
+        MLIPMap[model] = getattr(module, metadata["class"])
+    except ModuleNotFoundError as e:
+        print(e)
+        continue
+
+MLIPEnum = Enum("MLIPEnum", MLIPMap)
 
 class MLIP(
     nn.Module,

mlip_arena/models/externals/alignn.py

@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+from alignn.ff.ff import AlignnAtomwiseCalculator, default_path
+
+from mlip_arena.models.utils import get_freer_device
+
+
+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)

mlip_arena/models/externals/chgnet.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from typing import Literal
+
+from ase import Atoms
+from chgnet.model.dynamics import CHGNetCalculator
+from chgnet.model.model import CHGNet as CHGNetModel
+
+from mlip_arena.models.utils import get_freer_device
+
+
+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)

mlip_arena/models/externals/equiformer.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import yaml
+from ase import Atoms
+from fairchem.core import OCPCalculator
+
+with open(Path(__file__).parents[1] / "registry.yaml", encoding="utf-8") as f:
+    REGISTRY = yaml.safe_load(f)
+
+
+class EquiformerV2(OCPCalculator):
+    def __init__(
+        self,
+        checkpoint=REGISTRY["EquiformerV2(OC22)"]["checkpoint"],
+        # 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=REGISTRY["EquiformerV2(OC22)"]["checkpoint"],
+        # 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,
+        )

mlip_arena/models/externals/escn.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import yaml
+from ase import Atoms
+from fairchem.core import OCPCalculator
+
+with open(Path(__file__).parents[1] / "registry.yaml", encoding="utf-8") as f:
+    REGISTRY = yaml.safe_load(f)
+
+class eSCN(OCPCalculator):
+    def __init__(
+        self,
+        checkpoint=REGISTRY["eSCN(OC20)"]["checkpoint"], # "eSCN-L6-M3-Lay20-S2EF-OC20-All+MD"
+        # 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(),
+        )

mlip_arena/models/externals/fairchem.py

@@ -0,0 +1,124 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import yaml
+from ase import Atoms
+from fairchem.core import OCPCalculator
+from huggingface_hub import hf_hub_download
+
+with open(Path(__file__).parents[1] / "registry.yaml", encoding="utf-8") as f:
+    REGISTRY = yaml.safe_load(f)
+
+class eqV2(OCPCalculator):
+    def __init__(
+        self,
+        checkpoint=REGISTRY["eqV2(OMat)"]["checkpoint"],
+        cache_dir=None,
+        cpu=False, # TODO: cannot assign device
+        seed=0,
+        **kwargs,
+    ) -> None:
+        """
+        Initialize an eqV2 calculator.
+
+        Parameters
+        ----------
+        checkpoint : str, default="eqV2_86M_omat_mp_salex.pt"
+            The name of the eqV2 checkpoint to use.
+        local_cache : str, default="/tmp/ocp/"
+            The directory to store the downloaded checkpoint.
+        cpu : bool, default=False
+            Whether to run the model on CPU or GPU.
+        seed : int, default=0
+            The random seed for the model.
+
+        Other Parameters
+        ----------------
+        **kwargs
+            Any additional keyword arguments are passed to the superclass.
+        """
+
+        # https://huggingface.co/fairchem/OMAT24/resolve/main/eqV2_86M_omat_mp_salex.pt
+
+        checkpoint_path = hf_hub_download(
+            "fairchem/OMAT24",
+            filename=checkpoint,
+            revision="bf92f9671cb9d5b5c77ecb4aa8b317ff10b882ce",
+            cache_dir=cache_dir
+        )
+        super().__init__(
+            checkpoint_path=checkpoint_path,
+            cpu=cpu,
+            seed=seed,
+            **kwargs,
+        )
+
+class EquiformerV2(OCPCalculator):
+    def __init__(
+        self,
+        checkpoint=REGISTRY["EquiformerV2(OC22)"]["checkpoint"],
+        # 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=REGISTRY["EquiformerV2(OC22)"]["checkpoint"],
+        # 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(),
+        )

mlip_arena/models/externals/mace-mp.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+from mace.calculators import MACECalculator
+
+from mlip_arena.models.utils import get_freer_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
+        )

mlip_arena/models/externals/mace-off.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+from mace.calculators import MACECalculator
+
+from mlip_arena.models.utils import get_freer_device
+
+
+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
+        )

mlip_arena/models/externals/matgl.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+import matgl
+import torch
+from matgl.ext.ase import PESCalculator
+
+
+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)

mlip_arena/models/externals/orb.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import requests
+from orb_models.forcefield import pretrained
+from orb_models.forcefield.calculator import ORBCalculator
+
+from mlip_arena.models.utils import get_freer_device
+
+
+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/externals/sevennet.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+from sevenn.sevennet_calculator import SevenNetCalculator
+
+from mlip_arena.models.utils import get_freer_device
+
+
+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)

mlip_arena/models/registry.yaml

@@ -8,7 +8,7 @@ MACE-MP(M):
   last-update: 2024-03-25T14:30:00
   datetime: 2024-03-25T14:30:00 # TODO: Fake datetime
   datasets:
-    - atomind/mptrj # TODO: fake HF dataset repo
+    - MPTrj # TODO: fake HF dataset repo
   cpu-tasks:
     - alexandria
     - qmof
@@ -33,7 +33,7 @@ CHGNet:
   last-update: 2024-07-08T00:00:00
   datetime: 2024-07-08T00:00:00
   datasets:
-    - atomind/mptrj
+    - MPTrj
   gpu-tasks:
     - homonuclear-diatomics
     - stability
@@ -48,14 +48,14 @@ CHGNet:
 M3GNet:
   module: externals
   class: M3GNet
-  family: m3gnet
+  family: matgl
   package: matgl==1.1.2
   checkpoint:
   username: cyrusyc
   last-update: 2024-07-08T00:00:00
   datetime: 2024-07-08T00:00:00
   datasets:
-    - atomind/mptrj
+    - MPF
   gpu-tasks:
     - homonuclear-diatomics
     - combustion
@@ -76,8 +76,8 @@ ORB:
   last-update: 2024-03-25T14:30:00
   datetime: 2024-03-25T14:30:00 # TODO: Fake datetime
   datasets:
-    - atomind/mptrj # TODO: fake HF dataset repo
-    - atomind/alexandria
+    - MPTrj # TODO: fake HF dataset repo
+    - Alexandria
   cpu-tasks:
     - alexandria
     - qmof
@@ -102,7 +102,7 @@ SevenNet:
   last-update: 2024-03-25T14:30:00
   datetime: 2024-03-25T14:30:00 # TODO: Fake datetime
   datasets:
-    - atomind/mptrj # TODO: fake HF dataset repo
+    - MPTrj # TODO: fake HF dataset repo
   cpu-tasks:
     - alexandria
     - qmof
@@ -117,17 +117,39 @@ SevenNet:
   nvt: true
   npt: true
 
+eqV2(OMat):
+  module: externals
+  class: eqV2
+  family: fairchem
+  package: fairchem-core==1.1.0
+  checkpoint: eqV2_86M_omat_mp_salex.pt
+  username: fairchem # HF handle
+  last-update: 2024-10-18T00:00:00
+  datetime: 2024-10-18T00:00:00
+  datasets:
+    - OMat
+    - MPTrj
+    - Alexandria
+  gpu-tasks:
+    - homonuclear-diatomics
+  prediction: EFS
+  nvt: true
+  npt: true
+  github: https://github.com/FAIR-Chem/fairchem
+  doi: https://arxiv.org/abs/2410.12771
+
+
 EquiformerV2(OC22):
   module: externals
   class: EquiformerV2
   family: equiformer
-  package: fairchem-core==1.0.0
+  package: fairchem-core==1.1.0
   checkpoint: EquiformerV2-lE4-lF100-S2EFS-OC22
   username: cyrusyc
   last-update: 2024-07-08T00:00:00
   datetime: 2024-07-08T00:00:00
   datasets:
-    - ocp
+    - OC22
   gpu-tasks:
     - homonuclear-diatomics
     - combustion
@@ -142,13 +164,13 @@ EquiformerV2(OC20):
   module: externals
   class: EquiformerV2OC20
   family: equiformer
-  package: fairchem-core==1.0.0
+  package: fairchem-core==1.1.0
   checkpoint: EquiformerV2-31M-S2EF-OC20-All+MD
   username: cyrusyc
   last-update: 2024-07-08T00:00:00
   datetime: 2024-07-08T00:00:00
   datasets:
-    - ocp
+    - OC20
   gpu-tasks:
     - homonuclear-diatomics
     - combustion
@@ -163,13 +185,13 @@ eSCN(OC20):
   module: externals
   class: eSCN
   family: escn
-  package: fairchem-core==1.0.0
+  package: fairchem-core==1.1.0
   checkpoint: eSCN-L6-M3-Lay20-S2EF-OC20-All+MD
   username: cyrusyc
   last-update: 2024-07-08T00:00:00
   datetime: 2024-07-08T00:00:00
   datasets:
-    - ocp
+    - OC20
   gpu-tasks:
     - homonuclear-diatomics
     - combustion
@@ -188,7 +210,7 @@ MACE-OFF(M):
   last-update: 2024-03-25T14:30:00
   datetime: 2024-03-25T14:30:00 # TODO: Fake datetime
   datasets:
-    - atomind/mptrj # TODO: fake HF dataset repo
+    - SPICE # TODO: fake HF dataset repo
   cpu-tasks:
     - alexandria
     - qmof
@@ -211,7 +233,7 @@ ALIGNN:
   last-update: 2024-07-08T00:00:00
   datetime: 2024-07-08T00:00:00
   datasets:
-    - atomind/mptrj
+    - MP22
   gpu-tasks:
     - homonuclear-diatomics
     # - combustion

mlip_arena/models/utils.py

@@ -1,20 +1,7 @@
 """Utility functions for MLIP models."""
 
-import importlib
-from enum import Enum
-
 import torch
 
-from mlip_arena.models import REGISTRY
-
-MLIPMap = {
-    model: getattr(
-        importlib.import_module(f"{__package__}.{metadata['module']}"), metadata["class"],
-    )
-    for model, metadata in REGISTRY.items()
-}
-MLIPEnum = Enum("MLIPEnum", MLIPMap)
-
 
 def get_freer_device() -> torch.device:
     """Get the GPU with the most free memory, or use MPS if available.
@@ -47,4 +34,4 @@ def get_freer_device() -> torch.device:
         print("No GPU or MPS available. Using CPU.")
         device = torch.device("cpu")
 
-    return device
+    return device

mlip_arena/tasks/diatomics/fairchem/homonuclear-diatomics.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:16144d6e7ff5b05d805f37cef43c03ccbc1f27787242da1a12455f62708069a8
+size 2132550

mlip_arena/tasks/diatomics/matgl/homonuclear-diatomics.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c257248d2f62bfd09eb334c49405287bf5bc3bf14cfc6ad0a5890425f559f91c
+size 1854330

mlip_arena/tasks/eos/run.py

@@ -18,7 +18,7 @@ from prefect.futures import wait
 from prefect.runtime import flow_run, task_run
 from pymatgen.analysis.eos import BirchMurnaghan
 
-from mlip_arena.models.utils import MLIPEnum
+from mlip_arena.models import MLIPEnum
 from mlip_arena.tasks.optimize import run as OPT
 
 if TYPE_CHECKING:

mlip_arena/tasks/md.py

@@ -84,7 +84,8 @@ from scipy.linalg import schur
 from torch_dftd.torch_dftd3_calculator import TorchDFTD3Calculator
 from tqdm.auto import tqdm
 
-from mlip_arena.models.utils import MLIPEnum, get_freer_device
+from mlip_arena.models import MLIPEnum
+from mlip_arena.models.utils import get_freer_device
 
 _valid_dynamics: dict[str, tuple[str, ...]] = {
     "nve": ("velocityverlet",),

mlip_arena/tasks/optimize.py

@@ -17,7 +17,8 @@ from prefect import task
 from prefect.tasks import task_input_hash
 from torch_dftd.torch_dftd3_calculator import TorchDFTD3Calculator
 
-from mlip_arena.models.utils import MLIPEnum, get_freer_device
+from mlip_arena.models import MLIPEnum
+from mlip_arena.models.utils import get_freer_device
 
 _valid_filters: dict[str, Filter] = {
     "Filter": Filter,

mlip_arena/tasks/run.py

@@ -1,318 +0,0 @@
-from __future__ import annotations
-
-from datetime import datetime, timedelta
-from pathlib import Path
-from typing import Literal, Sequence, Tuple
-
-import numpy as np
-from ase import Atoms, units
-from ase.calculators.calculator import Calculator
-from ase.calculators.mixing import SumCalculator
-from ase.io import read
-from ase.io.trajectory import Trajectory
-from ase.md.andersen import Andersen
-from ase.md.langevin import Langevin
-from ase.md.md import MolecularDynamics
-from ase.md.npt import NPT
-from ase.md.nptberendsen import NPTBerendsen
-from ase.md.nvtberendsen import NVTBerendsen
-from ase.md.velocitydistribution import (
-    MaxwellBoltzmannDistribution,
-    Stationary,
-    ZeroRotation,
-)
-from ase.md.verlet import VelocityVerlet
-from prefect import task
-from prefect.tasks import task_input_hash
-from scipy.interpolate import interp1d
-from scipy.linalg import schur
-from torch_dftd.torch_dftd3_calculator import TorchDFTD3Calculator
-from tqdm.auto import tqdm
-
-from mlip_arena.models.utils import MLIPEnum, get_freer_device
-
-# from mlip_arena.models.utils import EXTMLIPEnum, MLIPMap, external_ase_calculator
-
-_valid_dynamics: dict[str, tuple[str, ...]] = {
-    "nve": ("velocityverlet",),
-    "nvt": ("nose-hoover", "langevin", "andersen", "berendsen"),
-    "npt": ("nose-hoover", "berendsen"),
-}
-
-_preset_dynamics: dict = {
-    "nve_velocityverlet": VelocityVerlet,
-    "nvt_andersen": Andersen,
-    "nvt_berendsen": NVTBerendsen,
-    "nvt_langevin": Langevin,
-    "nvt_nose-hoover": NPT,
-    "npt_berendsen": NPTBerendsen,
-    "npt_nose-hoover": NPT,
-}
-
-
-def _interpolate_quantity(values: Sequence | np.ndarray, n_pts: int) -> np.ndarray:
-    """Interpolate temperature / pressure on a schedule."""
-    n_vals = len(values)
-    return np.interp(
-        np.linspace(0, n_vals - 1, n_pts + 1),
-        np.linspace(0, n_vals - 1, n_vals),
-        values,
-    )
-
-
-def _get_ensemble_schedule(
-    ensemble: Literal["nve", "nvt", "npt"] = "nvt",
-    n_steps: int = 1000,
-    temperature: float | Sequence | np.ndarray | None = 300.0,
-    pressure: float | Sequence | np.ndarray | None = None,
-) -> Tuple[np.ndarray, np.ndarray]:
-    if ensemble == "nve":
-        # Disable thermostat and barostat
-        temperature = np.nan
-        pressure = np.nan
-        t_schedule = np.full(n_steps + 1, temperature)
-        p_schedule = np.full(n_steps + 1, pressure)
-        return t_schedule, p_schedule
-
-    if isinstance(temperature, Sequence) or (
-        isinstance(temperature, np.ndarray) and temperature.ndim == 1
-    ):
-        t_schedule = _interpolate_quantity(temperature, n_steps)
-    # NOTE: In ASE Langevin dynamics, the temperature are normally
-    # scalars, but in principle one quantity per atom could be specified by giving
-    # an array. This is not implemented yet here.
-    else:
-        t_schedule = np.full(n_steps + 1, temperature)
-
-    if ensemble == "nvt":
-        pressure = np.nan
-        p_schedule = np.full(n_steps + 1, pressure)
-        return t_schedule, p_schedule
-
-    if isinstance(pressure, Sequence) or (
-        isinstance(pressure, np.ndarray) and pressure.ndim == 1
-    ):
-        p_schedule = _interpolate_quantity(pressure, n_steps)
-    elif isinstance(pressure, np.ndarray) and pressure.ndim == 4:
-        p_schedule = interp1d(np.arange(n_steps + 1), pressure, kind="linear")
-        assert isinstance(p_schedule, np.ndarray)
-    else:
-        p_schedule = np.full(n_steps + 1, pressure)
-
-    return t_schedule, p_schedule
-
-
-def _get_ensemble_defaults(
-    ensemble: Literal["nve", "nvt", "npt"],
-    dynamics: str | MolecularDynamics,
-    t_schedule: np.ndarray,
-    p_schedule: np.ndarray,
-    ase_md_kwargs: dict | None = None,
-) -> dict:
-    """Update ASE MD kwargs"""
-    ase_md_kwargs = ase_md_kwargs or {}
-
-    if ensemble == "nve":
-        ase_md_kwargs.pop("temperature", None)
-        ase_md_kwargs.pop("temperature_K", None)
-        ase_md_kwargs.pop("externalstress", None)
-    elif ensemble == "nvt":
-        ase_md_kwargs["temperature_K"] = t_schedule[0]
-        ase_md_kwargs.pop("externalstress", None)
-    elif ensemble == "npt":
-        ase_md_kwargs["temperature_K"] = t_schedule[0]
-        ase_md_kwargs["externalstress"] = p_schedule[0]  # * 1e3 * units.bar
-
-    if isinstance(dynamics, str) and dynamics.lower() == "langevin":
-        ase_md_kwargs["friction"] = ase_md_kwargs.get(
-            "friction",
-            10.0 * 1e-3 / units.fs,  # Same default as in VASP: 10 ps^-1
-        )
-
-    return ase_md_kwargs
-
-
-@task(cache_key_fn=task_input_hash, cache_expiration=timedelta(days=1))
-def md(
-    atoms: Atoms,
-    calculator_name: str | MLIPEnum,
-    calculator_kwargs: dict | None,
-    dispersion: str | None = None,
-    dispersion_kwargs: dict | None = None,
-    device: str | None = None,
-    ensemble: Literal["nve", "nvt", "npt"] = "nvt",
-    dynamics: str | MolecularDynamics = "langevin",
-    time_step: float | None = None,
-    total_time: float = 1000,
-    temperature: float | Sequence | np.ndarray | None = 300.0,
-    pressure: float | Sequence | np.ndarray | None = None,
-    ase_md_kwargs: dict | None = None,
-    md_velocity_seed: int | None = None,
-    zero_linear_momentum: bool = True,
-    zero_angular_momentum: bool = True,
-    traj_file: str | Path | None = None,
-    traj_interval: int = 1,
-    restart: bool = True,
-):
-    device = device or str(get_freer_device())
-
-    print(f"Using device: {device}")
-
-    calculator_kwargs = calculator_kwargs or {}
-
-    if isinstance(calculator_name, MLIPEnum) and calculator_name in MLIPEnum:
-        assert issubclass(calculator_name.value, Calculator)
-        calc = calculator_name.value(**calculator_kwargs)
-    elif (
-        isinstance(calculator_name, str) and calculator_name in MLIPEnum._member_names_
-    ):
-        calc = MLIPEnum[calculator_name].value(**calculator_kwargs)
-    else:
-        raise ValueError(f"Invalid calculator: {calculator_name}")
-
-    print(f"Using calculator: {calc}")
-
-    dispersion_kwargs = dispersion_kwargs or {}
-
-    dispersion_kwargs.update({"device": device})
-
-    if dispersion is not None:
-        disp_calc = TorchDFTD3Calculator(
-            **dispersion_kwargs,
-        )
-        calc = SumCalculator([calc, disp_calc])
-
-        print(f"Using dispersion: {dispersion}")
-
-    atoms.calc = calc
-
-    if time_step is None:
-        # If a structure contains an isotope of hydrogen, set default `time_step`
-        # to 0.5 fs, and 2 fs otherwise.
-        has_h_isotope = "H" in atoms.get_chemical_symbols()
-        time_step = 0.5 if has_h_isotope else 2.0
-
-    n_steps = int(total_time / time_step)
-    target_steps = n_steps
-
-    t_schedule, p_schedule = _get_ensemble_schedule(
-        ensemble=ensemble,
-        n_steps=n_steps,
-        temperature=temperature,
-        pressure=pressure,
-    )
-
-    ase_md_kwargs = _get_ensemble_defaults(
-        ensemble=ensemble,
-        dynamics=dynamics,
-        t_schedule=t_schedule,
-        p_schedule=p_schedule,
-        ase_md_kwargs=ase_md_kwargs,
-    )
-
-    if isinstance(dynamics, str):
-        # Use known dynamics if `self.dynamics` is a str
-        dynamics = dynamics.lower()
-        if dynamics not in _valid_dynamics[ensemble]:
-            raise ValueError(
-                f"{dynamics} thermostat not available for {ensemble}."
-                f"Available {ensemble} thermostats are:"
-                " ".join(_valid_dynamics[ensemble])
-            )
-        if ensemble == "nve":
-            dynamics = "velocityverlet"
-        md_class = _preset_dynamics[f"{ensemble}_{dynamics}"]
-    elif dynamics is MolecularDynamics:
-        md_class = dynamics
-    else:
-        raise ValueError(f"Invalid dynamics: {dynamics}")
-
-    if md_class is NPT:
-        #  Note that until md_func is instantiated, isinstance(md_func,NPT) is False
-        # ASE NPT implementation requires upper triangular cell
-        u, _ = schur(atoms.get_cell(complete=True), output="complex")
-        atoms.set_cell(u.real, scale_atoms=True)
-
-    last_step = 0
-
-    if traj_file is not None:
-        traj_file = Path(traj_file)
-        traj_file.parent.mkdir(parents=True, exist_ok=True)
-
-        if restart and traj_file.exists():
-            try:
-                traj = read(traj_file, index=":")
-                last_atoms = traj[-1]
-                assert isinstance(last_atoms, Atoms)
-                last_step = last_atoms.info.get("step", len(traj) * traj_interval)
-                n_steps -= last_step
-                traj = Trajectory(traj_file, "a", atoms)
-                atoms.set_positions(last_atoms.get_positions())
-                atoms.set_momenta(last_atoms.get_momenta())
-            except Exception:
-                traj = Trajectory(traj_file, "w", atoms)
-
-                if not np.isnan(t_schedule).any():
-                    MaxwellBoltzmannDistribution(
-                        atoms=atoms,
-                        temperature_K=t_schedule[last_step],
-                        rng=np.random.default_rng(seed=md_velocity_seed),
-                    )
-
-                if zero_linear_momentum:
-                    Stationary(atoms)
-                if zero_angular_momentum:
-                    ZeroRotation(atoms)
-        else:
-            traj = Trajectory(traj_file, "w", atoms)
-
-            if not np.isnan(t_schedule).any():
-                MaxwellBoltzmannDistribution(
-                    atoms=atoms,
-                    temperature_K=t_schedule[last_step],
-                    rng=np.random.default_rng(seed=md_velocity_seed),
-                )
-
-            if zero_linear_momentum:
-                Stationary(atoms)
-            if zero_angular_momentum:
-                ZeroRotation(atoms)
-
-    md_runner = md_class(
-        atoms=atoms,
-        timestep=time_step * units.fs,
-        **ase_md_kwargs,
-    )
-
-    if traj_file is not None:
-        md_runner.attach(traj.write, interval=traj_interval)
-
-    with tqdm(total=n_steps) as pbar:
-
-        def _callback(dyn: MolecularDynamics = md_runner) -> None:
-            step = last_step + dyn.nsteps
-            dyn.atoms.info["restart"] = last_step
-            dyn.atoms.info["datetime"] = datetime.now()
-            dyn.atoms.info["step"] = step
-            dyn.atoms.info["target_steps"] = target_steps
-            if ensemble == "nve":
-                return
-            dyn.set_temperature(temperature_K=t_schedule[step])
-            if ensemble == "nvt":
-                return
-            dyn.set_stress(p_schedule[step] * 1e3 * units.bar)
-            pbar.update()
-
-        md_runner.attach(_callback, interval=1)
-
-        start_time = datetime.now()
-        md_runner.run(steps=n_steps)
-        end_time = datetime.now()
-
-        traj.close()
-
-    return {
-        "atoms": atoms,
-        "runtime": end_time - start_time,
-        "n_steps": n_steps,
-    }

pyproject.toml

@@ -44,11 +44,12 @@ run = [
     "dgl==2.4.0",
     "mace-torch==0.3.4",
     "chgnet==0.3.8",
-    "fairchem-core==1.1.0",
+    "fairchem-core==1.2.0",
     "sevenn==0.9.3.post1",
     "orb-models==0.3.1",
     "alignn==2024.5.27",
-    "prefect>=3.0.4"
+    "prefect>=3.0.4",
+    "prefect-dask"
 ]
 app = [
     "streamlit==1.38.0",
@@ -62,7 +63,7 @@ test = [
     "matgl==1.1.2",
     "dgl==2.4.0",
     "chgnet==0.3.8",
-    "fairchem-core==1.1.0",
+    "fairchem-core==1.2.0",
     "sevenn==0.9.3.post1",
     "orb-models==0.3.1",
     "alignn==2024.5.27",

serve/leaderboard.py

@@ -25,6 +25,7 @@ table = pd.DataFrame(
         "Prediction",
         "NVT",
         "NPT",
+        "Training Set",
         "Code",
         "Paper",
         "First Release",
@@ -42,6 +43,7 @@ for model in MODELS:
         "Prediction": metadata.get("prediction", None),
         "NVT": "✅" if metadata.get("nvt", False) else "❌",
         "NPT": "✅" if metadata.get("npt", False) else "❌",
+        "Training Set": metadata.get("datasets", []),
         "Code": metadata.get("github", None) if metadata else None,
         "Paper": metadata.get("doi", None) if metadata else None,
         "First Release": metadata.get("date", None),
@@ -50,7 +52,6 @@ for model in MODELS:
 
 table.set_index("Model", inplace=True)
 
-
 s = table.style.background_gradient(
     cmap="PuRd", subset=["Element Coverage"], vmin=0, vmax=120
 )

serve/ranks/homonuclear-diatomics.py

@@ -72,6 +72,8 @@ 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["Rank"] += 1
+
 table.sort_values("Rank", ascending=True, inplace=True)
 
 table["Rank aggr."] = table["Rank"]

serve/tasks/homonuclear-diatomics.py

@@ -30,7 +30,7 @@ valid_models = [
 mlip_methods = container.multiselect(
     "MLIPs",
     valid_models,
-    ["EquiformerV2(OC22)", "CHGNet", "M3GNet", "SevenNet", "MACE-MP(M)", "ORB"],
+    ["EquiformerV2(OC22)", "CHGNet", "M3GNet", "SevenNet", "MACE-MP(M)", "ORB", "eqV2(OMat)"],
 )
 dft_methods = container.multiselect("DFT Methods", ["GPAW"], [])
 
@@ -139,6 +139,7 @@ def get_plots(df, energy_plot: bool, force_plot: bool, method_color_mapping: dic
                     ys = es
 
                 elo = min(elo, max(ys.min() * 1.2, -15), -1)
+                # elo = min(elo, ys.min()*1.2, -1)
 
                 fig.add_trace(
                     go.Scatter(
@@ -202,7 +203,7 @@ def get_plots(df, energy_plot: bool, force_plot: bool, method_color_mapping: dic
                 yaxis=dict(
                     title=dict(text="Energy [eV]"),
                     side="left",
-                    range=[elo, 1.5 * (abs(elo))],
+                    range=[elo, 2.0 * (abs(elo))],
                 )
             )
 

tests/test_eos.py

@@ -1,34 +1,41 @@
 import pytest
 from ase.build import bulk
+from prefect.testing.utilities import prefect_test_harness
 
-from mlip_arena.models.utils import MLIPEnum
+from mlip_arena.models import MLIPEnum
 from mlip_arena.tasks.eos.run import fit as EOS
+import sys
 
 atoms = bulk("Cu", "fcc", a=3.6)
 
+# @pytest.fixture(autouse=True, scope="session")
+# def prefect_test_fixture():
+#     with prefect_test_harness():
+#         yield
 
+@pytest.mark.skipif(sys.version_info[:2] != (3,11), reason="avoid prefect race condition on concurrent tasks")
 @pytest.mark.parametrize("model", [MLIPEnum["MACE-MP(M)"]])
 def test_eos(model: MLIPEnum):
     """
     Test EOS prefect workflow with a simple cubic lattice.
     """
 
-    result = EOS(
-        atoms=atoms,
-        calculator_name=model.name,
-        calculator_kwargs={},
-        device=None,
-        optimizer="BFGSLineSearch",
-        optimizer_kwargs=None,
-        filter="FrechetCell",
-        filter_kwargs=None,
-        criterion=dict(
-            fmax=0.1,
-        ),
-        max_abs_strain=0.1,
-        npoints=6,
-    )
+    with prefect_test_harness():
 
-    print(result)
+        result = EOS(
+            atoms=atoms,
+            calculator_name=model.name,
+            calculator_kwargs={},
+            device=None,
+            optimizer="BFGSLineSearch",
+            optimizer_kwargs=None,
+            filter="FrechetCell",
+            filter_kwargs=None,
+            criterion=dict(
+                fmax=0.1,
+            ),
+            max_abs_strain=0.1,
+            npoints=6,
+        )
 
-    assert isinstance(result["K"], float)
+        assert isinstance(result["K"], float)

tests/test_external_calculators.py

@@ -1,7 +1,7 @@
 import pytest
 from ase import Atoms
 
-from mlip_arena.models.utils import MLIPEnum
+from mlip_arena.models import MLIPEnum
 
 
 @pytest.mark.parametrize("model", MLIPEnum)