refactor external calculators; better handle devices

mlip_arena/models/__init__.py

@@ -18,37 +18,29 @@ class MLIP(
     PyTorchModelHubMixin,
     tags=["atomistic-simulation", "MLIP"],
 ):
-    def __init__(self, *args, **kwargs) -> None:
-        super().__init__(*args, **kwargs)
-
-
-class ModuleMLIP(MLIP):
-    def __init__(self, model: nn.Module, *args, **kwargs) -> None:
-        super().__init__(*args, **kwargs)
-        self.add_module("model", model)
+    def __init__(self, model: nn.Module) -> None:
+        super().__init__()
+        self.model = model
 
     def forward(self, x):
-        print("Forwarding...")
-        out = self.model(x)
-        print("Forwarded!")
-        return out
-
+        return self.model(x)
 
-class MLIPCalculator(Calculator):
+class MLIPCalculator(MLIP, Calculator):
     name: str
-    # device: torch.device
-    # model: MLIP
     implemented_properties: list[str] = ["energy", "forces", "stress"]
 
     def __init__(
         self,
+        model,
         # ASE Calculator
         restart=None,
         atoms=None,
         directory=".",
-        **kwargs,
+        calculator_kwargs: dict = {},
     ):
-        super().__init__(restart=restart, atoms=atoms, directory=directory, **kwargs)
+        MLIP.__init__(self, model=model)  # Initialize MLIP part
+        Calculator.__init__(self, restart=restart, atoms=atoms, directory=directory, **calculator_kwargs)  # Initialize ASE Calculator part
+        # Additional initialization if needed
         # self.name: str = self.__class__.__name__
         # self.device = device or torch.device(
         #     "cuda" if torch.cuda.is_available() else "cpu"

mlip_arena/models/chgnet.py

@@ -7,10 +7,12 @@ 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, ModuleMLIP
+from mlip_arena.models import MLIP, MLIPCalculator
 
+# TODO: WIP
 
-class CHGNetCalculator(MLIPCalculator):
+
+class CHGNet(MLIPCalculator):
     def __init__(
         self,
         device: torch.device | None = None,
@@ -19,23 +21,15 @@ class CHGNetCalculator(MLIPCalculator):
         directory=".",
         **kwargs,
     ):
-        super().__init__(restart=restart, atoms=atoms, directory=directory, **kwargs)
-
-        self.name: str = self.__class__.__name__
-
-        fpath = hf_hub_download(
-            repo_id="cyrusyc/mace-universal",
-            subfolder="pretrained",
-            filename="2023-12-12-mace-128-L1_epoch-199.model",
-            revision="main",
-        )
-
         self.device = device or torch.device(
             "cuda" if torch.cuda.is_available() else "cpu"
         )
 
-        self.model = torch.load(fpath, map_location=self.device)
+        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(

mlip_arena/models/externals.py

@@ -0,0 +1,165 @@
+import os
+import urllib
+from typing import Literal
+
+import torch
+from alignn.ff.ff import AlignnAtomwiseCalculator, get_figshare_model_ff
+from ase import Atoms
+from chgnet.model.dynamics import CHGNetCalculator
+from chgnet.model.model import CHGNet
+from fairchem.core import OCPCalculator
+from mace.calculators import MACECalculator
+
+
+# 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, device=None, default_dtype="float32", **kwargs):
+        checkpoint_url = "http://tinyurl.com/5yyxdm76"
+        cache_dir = os.path.expanduser("~/.cache/mace")
+        checkpoint_url_name = "".join(
+            c for c in os.path.basename(checkpoint_url) if c.isalnum() or c in "_"
+        )
+        cached_model_path = f"{cache_dir}/{checkpoint_url_name}"
+        if not os.path.isfile(cached_model_path):
+            os.makedirs(cache_dir, exist_ok=True)
+            # download and save to disk
+            print(f"Downloading MACE model from {checkpoint_url!r}")
+            _, http_msg = urllib.request.urlretrieve(checkpoint_url, cached_model_path)
+            if "Content-Type: text/html" in http_msg:
+                raise RuntimeError(
+                    f"Model download failed, please check the URL {checkpoint_url}"
+                )
+            print(f"Cached MACE model to {cached_model_path}")
+        model = cached_model_path
+        msg = f"Using Materials Project MACE for MACECalculator with {model}"
+        print(msg)
+
+        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,
+        model: CHGNet | None = None,
+        use_device: str | None = None,
+        stress_weight: float | None = 1 / 160.21766208,
+        on_isolated_atoms: Literal["ignore", "warn", "error"] = "warn",
+        **kwargs,
+    ) -> None:
+        use_device = use_device or str(get_freer_device())
+        super().__init__(
+            model=model,
+            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 EquiformerV2(OCPCalculator):
+    def __init__(
+        self,
+        model_name="EquiformerV2-lE4-lF100-S2EFS-OC22",
+        local_cache="/tmp/ocp/",
+        cpu=False,
+        seed=0,
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            model_name=model_name,
+            local_cache=local_cache,
+            cpu=cpu,
+            seed=0,
+            **kwargs,
+        )
+
+    def calculate(self, atoms: Atoms, properties, system_changes) -> None:
+        super().calculate(atoms, properties, system_changes)
+
+        self.results.update(
+            force=atoms.get_forces(),
+        )
+
+
+class eSCN(OCPCalculator):
+    def __init__(
+        self,
+        model_name="eSCN-L6-M3-Lay20-S2EF-OC20-All+MD",
+        local_cache="/tmp/ocp/",
+        cpu=False,
+        seed=0,
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            model_name=model_name,
+            local_cache=local_cache,
+            cpu=cpu,
+            seed=0,
+            **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, dir_path: str = "/tmp/alignn/", device=None, **kwargs) -> None:
+        model_path = get_figshare_model_ff(dir_path=dir_path)
+        device = device or get_freer_device()
+        super().__init__(model_path=model_path, device=device, **kwargs)
+
+    def calculate(self, atoms, properties=None, system_changes=None):
+        super().calculate(atoms, properties, system_changes)

mlip_arena/models/mace.py

@@ -1,13 +1,10 @@
-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, ModuleMLIP
+from mlip_arena.models import MLIPCalculator
 
 
 class MACE_MP_Medium(MLIPCalculator):
@@ -19,9 +16,9 @@ class MACE_MP_Medium(MLIPCalculator):
         directory=".",
         **kwargs,
     ):
-        super().__init__(restart=restart, atoms=atoms, directory=directory, **kwargs)
-
-        self.name: str = self.__class__.__name__
+        self.device = device or torch.device(
+            "cuda" if torch.cuda.is_available() else "cpu"
+        )
 
         fpath = hf_hub_download(
             repo_id="cyrusyc/mace-universal",
@@ -30,23 +27,15 @@ class MACE_MP_Medium(MLIPCalculator):
             revision="main",
         )
 
-        self.device = device or torch.device(
-            "cuda" if torch.cuda.is_available() else "cpu"
-        )
+        model = torch.load(fpath, map_location=self.device)
 
-        self.model = torch.load(fpath, map_location=self.device)
+        super().__init__(
+            model=model, restart=restart, atoms=atoms, directory=directory, **kwargs
+        )
 
+        self.name: str = self.__class__.__name__
         self.implemented_properties = ["energy", "forces", "stress"]
 
-        # repo_id = f"atomind/{self.__class__.__name__}".lower().replace("_", "-")
-
-        # model = ModuleMLIP(model=model)
-        # model.save_pretrained(
-        #     self.__class__.__name__.lower().replace("_", "-"),
-        #     repo_id=repo_id,
-        #     push_to_hub=True,
-        # )
-
     def calculate(
         self, atoms: Atoms, properties: list[str], system_changes: list = all_changes
     ):

mlip_arena/models/registry.yaml

@@ -1,7 +1,7 @@
 
 
 MACE-MP(M):
-  module: mace
+  module: externals
   class: MACE_MP_Medium
   username: cyrusyc # HF username
   last-update: 2024-03-25T14:30:00
@@ -17,6 +17,40 @@ MACE-MP(M):
   doi: https://arxiv.org/abs/2401.00096
   date: 2023-12-29
 
+CHGNet:
+  module: externals
+  class: CHGNet
+  username: cyrusyc
+  last-update: 2024-07-08T00:00:00
+  datetime: 2024-07-08T00:00:00
+  datasets:
+    - atomind/mptrj
+  gpu-tasks:
+    - diatomics
+
+EquiformerV2(OC22):
+  module: externals
+  class: EquiformerV2
+  username: cyrusyc
+  last-update: 2024-07-08T00:00:00
+  datetime: 2024-07-08T00:00:00
+  datasets:
+    - ocp
+  gpu-tasks:
+    - diatomics
+
+eSCN(OC20):
+  module: externals
+  class: eSCN
+  username: cyrusyc
+  last-update: 2024-07-08T00:00:00
+  datetime: 2024-07-08T00:00:00
+  datasets:
+    - ocp
+  gpu-tasks:
+    - diatomics
+
+
 # CHGNet:
 #   module: chgnet
 #   username: cyrusyc

mlip_arena/models/utils.py

@@ -2,10 +2,8 @@
 
 import importlib
 from enum import Enum
-from typing import Any
 
 import torch
-from ase.calculators.calculator import Calculator
 
 from mlip_arena.models import REGISTRY
 
@@ -15,73 +13,112 @@ MLIPMap = {
     )
     for model, metadata in REGISTRY.items()
 }
+MLIPEnum = Enum("MLIPEnum", MLIPMap)
 
 
-class EXTMLIPEnum(Enum):
-    """Enumeration class for EXTMLIP models.
+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.
 
-    Attributes:
-        M3GNet (str): M3GNet model.
-        CHGNet (str): CHGNet model.
-        MACE (str): MACE model.
+        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")
 
-    M3GNet = "M3GNet"
-    CHGNet = "CHGNet"
-    MACE = "MACE"
+    return device
 
 
-def get_freer_device() -> torch.device:
-    """Get the GPU with the most free memory.
+# class EXTMLIPEnum(Enum):
+#     """Enumeration class for EXTMLIP models.
 
-    Returns:
-        torch.device: The selected GPU device.
+#     Attributes:
+#         M3GNet (str): M3GNet model.
+#         CHGNet (str): CHGNet model.
+#         MACE (str): MACE model.
+#     """
 
-    Raises:
-        ValueError: If no GPU is available.
-    """
-    device_count = torch.cuda.device_count()
-    if device_count == 0:
-        print("No GPU available. Using CPU.")
-        return torch.device("cpu")
+#     M3GNet = "M3GNet"
+#     CHGNet = "CHGNet"
+#     MACE = "MACE"
+#     Equiformer = "Equiformer"
 
-    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))
+# def get_freer_device() -> torch.device:
+#     """Get the GPU with the most free memory.
+
+#     Returns:
+#         torch.device: The selected GPU device.
+
+#     Raises:
+#         ValueError: If no GPU is available.
+#     """
+#     device_count = torch.cuda.device_count()
+#     if device_count == 0:
+#         print("No GPU available. Using CPU.")
+#         return torch.device("cpu")
+
+#     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))
+
+#     print(
+#         f"Selected GPU {free_gpu_index} with {mem_free[free_gpu_index] / 1024**2:.2f} MB free memory from {device_count} GPUs",
+#     )
+
+#     return torch.device(f"cuda:{free_gpu_index}")
+
 
-    print(
-        f"Selected GPU {free_gpu_index} with {mem_free[free_gpu_index] / 1024**2:.2f} MB free memory from {device_count} GPUs",
-    )
 
-    return torch.device(f"cuda:{free_gpu_index}")
+# def external_ase_calculator(name: EXTMLIPEnum, **kwargs: Any) -> Calculator:
+#     """Construct an ASE calculator from an external third-party MLIP packages"""
+#     calculator = None
+#     device = get_freer_device()
 
+#     if name == EXTMLIPEnum.MACE:
+#         from mace.calculators import mace_mp
 
-def external_ase_calculator(name: EXTMLIPEnum, **kwargs: Any) -> Calculator:
-    """Construct an ASE calculator from an external third-party MLIP packages"""
-    calculator = None
-    device = get_freer_device()
+#         calculator = mace_mp(device=str(device), **kwargs)
 
-    if name == EXTMLIPEnum.MACE:
-        from mace.calculators import mace_mp
+#     elif name == EXTMLIPEnum.CHGNet:
+#         from chgnet.model.dynamics import CHGNetCalculator
 
-        calculator = mace_mp(device=str(device), **kwargs)
+#         calculator = CHGNetCalculator(use_device=str(device), **kwargs)
 
-    elif name == EXTMLIPEnum.CHGNet:
-        from chgnet.model.dynamics import CHGNetCalculator
+#     elif name == EXTMLIPEnum.M3GNet:
+#         import matgl
+#         from matgl.ext.ase import PESCalculator
 
-        calculator = CHGNetCalculator(use_device=str(device), **kwargs)
+#         potential = matgl.load_model("M3GNet-MP-2021.2.8-PES")
+#         calculator = PESCalculator(potential, **kwargs)
 
-    elif name == EXTMLIPEnum.M3GNet:
-        import matgl
-        from matgl.ext.ase import PESCalculator
 
-        potential = matgl.load_model("M3GNet-MP-2021.2.8-PES")
-        calculator = PESCalculator(potential, **kwargs)
 
-    calculator.__setattr__("name", name.value)
+#     calculator.__setattr__("name", name.value)
 
-    return calculator
+#     return calculator

mlip_arena/tasks/stability/__init__.py

@@ -0,0 +1,3 @@
+
+from .run import md as MD
+

mlip_arena/tasks/stability/run.py

@@ -1,39 +1,141 @@
 from __future__ import annotations
 
-import datetime
-from datetime import datetime
+from datetime import datetime, timedelta
 from pathlib import Path
-from typing import Literal, Sequence
+from typing import Literal, Sequence, Tuple
 
 import numpy as np
-import torch
 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 EXTMLIPEnum, MLIPMap, external_ase_calculator
-from mlip_arena.tasks.utils import (
-    _get_ensemble_defaults,
-    _get_ensemble_schedule,
-    _preset_dynamics,
-    _valid_dynamics,
-)
+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 | EXTMLIPEnum,
+    calculator_name: str | MLIPEnum,
     calculator_kwargs: dict | None,
     dispersion: str | None = None,
     dispersion_kwargs: dict | None = None,
@@ -59,16 +161,19 @@ def md(
     # device: str | None = None,
     # dtype: str = "float64",
 ):
-    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+    device = device or str(get_freer_device())
 
     print(f"Using device: {device}")
 
     calculator_kwargs = calculator_kwargs or {}
 
-    if isinstance(calculator_name, EXTMLIPEnum) and calculator_name in EXTMLIPEnum:
-        calc = external_ase_calculator(calculator_name, **calculator_kwargs)
-    elif calculator_name in MLIPMap:
-        calc = MLIPMap[calculator_name](**calculator_kwargs)
+    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}")
 
@@ -171,12 +276,15 @@ def md(
     with tqdm(total=n_steps) as pbar:
 
         def _callback(dyn: MolecularDynamics = md_runner) -> None:
+            step = last_step + dyn.nsteps
+            dyn.atoms.info["datetime"] = datetime.now()
+            dyn.atoms.info["step"] = step
             if ensemble == "nve":
                 return
-            dyn.set_temperature(temperature_K=t_schedule[last_step + dyn.nsteps])
+            dyn.set_temperature(temperature_K=t_schedule[step])
             if ensemble == "nvt":
                 return
-            dyn.set_stress(p_schedule[last_step + dyn.nsteps] * 1e3 * units.bar)
+            dyn.set_stress(p_schedule[step] * 1e3 * units.bar)
             pbar.update()
 
         md_runner.attach(_callback, interval=1)
@@ -187,4 +295,4 @@ def md(
 
         traj.close()
 
-    return {"md_runtime": end_time - start_time}
+    return {"runtime": end_time - start_time, "n_steps": n_steps}

mlip_arena/tasks/utils.py

@@ -1,161 +0,0 @@
-import os, glob
-from pathlib import Path
-from ase.io import read, write
-from ase import units
-from ase import Atoms, units
-from ase.calculators.calculator import Calculator
-from ase.data import chemical_symbols
-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 dask.distributed import Client
-from dask_jobqueue import SLURMCluster
-from jobflow import Maker
-from prefect import flow, task
-from prefect.tasks import task_input_hash
-from prefect_dask import DaskTaskRunner
-from pymatgen.io.ase import AseAtomsAdaptor
-from scipy.interpolate import interp1d
-from scipy.linalg import schur
-
-from mlip_arena.models import MLIPCalculator
-from mlip_arena.models.utils import EXTMLIPEnum, MLIPMap, external_ase_calculator
-from torch_dftd.torch_dftd3_calculator import TorchDFTD3Calculator
-from mp_api.client import MPRester
-
-from fireworks import LaunchPad
-from atomate2.vasp.flows.core import RelaxBandStructureMaker
-from atomate2.vasp.flows.mp import MPGGADoubleRelaxStaticMaker
-from atomate2.vasp.powerups import add_metadata_to_flow
-from atomate2.forcefields.md import (
-    CHGNetMDMaker,
-    GAPMDMaker,
-    M3GNetMDMaker,
-    MACEMDMaker,
-    NequipMDMaker,
-)
-from atomate2.forcefields.utils import MLFF
-from pymatgen.io.ase import AseAtomsAdaptor
-from pymatgen.transformations.advanced_transformations import CubicSupercellTransformation
-from jobflow.managers.fireworks import flow_to_workflow
-from jobflow import run_locally, SETTINGS
-from tqdm.auto import tqdm
-
-from datetime import timedelta, datetime
-from typing import Literal, Sequence, Tuple
-
-import numpy as np
-import torch
-from pymatgen.core.structure import Structure
-
-from ase.calculators.mixing import SumCalculator
-from scipy.interpolate import interp1d
-
-from ase.io.trajectory import Trajectory
-
-
-_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
-

tests/hf_hub.ipynb

@@ -2,22 +2,192 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from huggingface_hub import hf_hub_download\n",
+    "from ase.calculators.calculator import Calculator\n",
+    "# from mlip_arena.models import MLIP, MLIPCalculator, ModuleMLIP\n",
+    "\n",
+    "from mlip_arena.models.externals import MACE_MP_Medium\n",
+    "\n",
+    "from mlip_arena.models.utils import MLIPMap, MLIPEnum"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "True"
+      ]
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "issubclass(MLIPEnum[\"MACE-MP(M)\"].value, Calculator)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
    "metadata": {},
    "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "True"
+      ]
+     },
+     "execution_count": 13,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "isinstance(MLIPEnum[\"MACE-MP(M)\"], MLIPEnum)# in MLIPEnum"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['MACE-MP(M)', 'CHGNet', 'EquiformerV2(OC22)', 'eSCN(OC20)']"
+      ]
+     },
+     "execution_count": 16,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "MLIPEnum._member_names_"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'MACE-MP(M)': mlip_arena.models.externals.MACE_MP_Medium,\n",
+       " 'CHGNet': mlip_arena.models.externals.CHGNet,\n",
+       " 'EquiformerV2(OC22)': mlip_arena.models.externals.EquiformerV2,\n",
+       " 'eSCN(OC20)': mlip_arena.models.externals.eSCN}"
+      ]
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "MLIPMap"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "MLIPEnum.MACE-MP(M)\n",
+      "MLIPEnum.CHGNet\n",
+      "MLIPEnum.EquiformerV2(OC22)\n",
+      "MLIPEnum.eSCN(OC20)\n"
+     ]
+    }
+   ],
+   "source": [
+    "for mlip in MLIPEnum:\n",
+    "    print(mlip)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "mlip_arena.models.externals.MACE_MP_Medium"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "MLIPMap['MACE-MP(M)']"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Using Materials Project MACE for MACECalculator with /global/homes/c/cyrusyc/.cache/mace/5yyxdm76\n",
+      "Selected GPU cuda:0 with 40338.06 MB free memory from 1 GPUs\n",
+      "Default dtype float32 does not match model dtype float64, converting models to float32.\n"
+     ]
+    }
+   ],
+   "source": [
+    "mace_mp = MACE_MP_Medium()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Select GPU cuda:0 with 40316.98 MB free memory from 1 GPUs\n",
+      "CHGNet v0.3.0 initialized with 412,525 parameters\n",
+      "CHGNet will run on cuda:0\n"
+     ]
+    },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
-      "  from .autonotebook import tqdm as notebook_tqdm\n"
+      "WARNING:root:Detected old config, converting to new format. Consider updating to avoid potential incompatibilities.\n",
+      "WARNING:root:Skipping scheduler setup. No training set found.\n"
      ]
     }
    ],
    "source": [
-    "import torch\n",
-    "from huggingface_hub import hf_hub_download\n",
-    "from mlip_arena.models import MLIP, MLIPCalculator, ModuleMLIP"
+    "from mlip_arena.models.externals import EquiformerV2, CHGNet\n",
+    "\n",
+    "chgnet = CHGNet()\n",
+    "\n",
+    "equiformer_v2 = EquiformerV2()\n"
    ]
   },
   {