Add eSEN and update dependency versions (#58)

* add prefect task flow script

* add esen diatomics

* separate matgl

* update installation script

* adjust test ci installation

* enforce `dgl==2.5.0`

* rollback torch to 2.4

* test first

* dgl 2.2.1

* dgl 2.2.1

* dgl 2.2.1

* no cache; cu124

* remove pin for dgl

* bump matgl version and update calculator api

* rollback orb to 0.4.0

* specify fairchem github repo as dep

* update uv version on ci

* fairchem

.github/README.md

@@ -34,28 +34,41 @@ pip install mlip-arena
 > [!CAUTION] 
 > We recommend clean build in a new virtual environment due to the compatibility issues between multiple popular MLIPs. We provide a single installation script using `uv` for minimal package conflicts and fast installation!
 
+> [!CAUTION]
+> To automatically download farichem OMat24 checkpoint, please make sure you have gained downloading access to their HuggingFace [***model repo***](https://huggingface.co/facebook/OMAT24) (not dataset repo), and login locally on your machine through `huggginface-cli login` (see [HF hub authentication](https://huggingface.co/docs/huggingface_hub/en/quick-start#authentication))
+
 **Linux**
 
 ```bash
 # (Optional) Install uv
 curl -LsSf https://astral.sh/uv/install.sh | sh
 source $HOME/.local/bin/env
+
+git clone https://github.com/atomind-ai/mlip-arena.git
+cd mlip-arena
+
 # One script uv pip installation
 bash scripts/install-linux.sh
 ```
 
-```bash
+<!-- ```bash
 # Or from command line
 git clone https://github.com/atomind-ai/mlip-arena.git
 cd mlip-arena
-pip install torch==2.2.0
-bash scripts/install-pyg.sh
-bash scripts/install-dgl.sh
+
+TORCH=2.4
+CUDA=cu124
+pip install torch==${TORCH}.0
+pip install torch-scatter torch-sparse -f https://data.pyg.org/whl/torch-${TORCH}.0+${CUDA}.html
+pip install dgl -f https://data.dgl.ai/wheels/torch-${TORCH}/${CUDA}/repo.html
+pip install -e .[fairchem] --no-deps
+pip install -e .[orb]
+pip install -e .[matgl]
 pip install -e .[test]
 pip install -e .[mace]
 # DeePMD
 DP_ENABLE_TENSORFLOW=0 pip install -e .[deepmd]
-```
+``` -->
 
 **Mac**
 

.github/workflows/test.yaml

@@ -24,8 +24,9 @@ jobs:
       uses: actions/checkout@v4
 
     - name: Install uv
-      uses: astral-sh/setup-uv@v4
+      uses: astral-sh/setup-uv@v6
       with:
+        version: "latest"
         enable-cache: true
         cache-dependency-glob: "pyproject.toml"
 
@@ -36,15 +37,7 @@ jobs:
 
     - name: Install dependencies
       run: |
-        TORCH=2.2
-        CUDA=cu121
-        uv pip install torch==${TORCH}.0
-        uv pip install torch-scatter -f https://data.pyg.org/whl/torch-${TORCH}.0+${CUDA}.html
-        uv pip install torch-sparse -f https://data.pyg.org/whl/torch-${TORCH}.0+${CUDA}.html
-        uv pip install dgl -f https://data.dgl.ai/wheels/torch-${TORCH}/${CUDA}/repo.html
-        uv pip install -e .[test]
-        uv pip install -e .[mace]
-        uv pip install -e .[deepmd]
+        bash scripts/install-linux.sh
 
     - name: List dependencies
       run: pip list

examples/mof/classification/classification.py

@@ -60,7 +60,7 @@ def save_result(
 # Orchestrate your awesome dask workflow runner
 
 nodes_per_alloc = 1
-gpus_per_alloc = 4
+gpus_per_alloc = 1
 ntasks = 1
 
 cluster_kwargs = dict(
@@ -68,7 +68,7 @@ cluster_kwargs = dict(
     memory="64 GB",
     shebang="#!/bin/bash",
     account="m3828",
-    walltime="01:30:00",
+    walltime="04:00:00",
     job_mem="0",
     job_script_prologue=[
         "source ~/.bashrc",
@@ -82,7 +82,7 @@ cluster_kwargs = dict(
         f"-N {nodes_per_alloc}",
         "-C gpu",
         f"-G {gpus_per_alloc}",
-        "--exclusive",
+        # "--exclusive",
     ],
 )
 
@@ -122,7 +122,7 @@ def run_all():
 
     for model, row in tqdm(itertools.product(MLIPEnum, load_row_from_df("input.pkl"))):
 
-        if model.name not in ["MACE-MPA", "MatterSim", "SevenNet", "M3GNet", "CHGNet", "ORBv2"]:
+        if model.name not in ["MACE-MPA", "MatterSim", "SevenNet", "M3GNet", "ORBv2"]:
             continue
 
         fpath = Path(f"{model.name}.pkl")

mlip_arena/models/__init__.py

@@ -42,7 +42,7 @@ for model, metadata in REGISTRY.items():
             f"{__package__}.{metadata['module']}.{metadata['family']}"
         )
         MLIPMap[model] = getattr(module, metadata["class"])
-    except (ModuleNotFoundError, AttributeError, ValueError) as e:
+    except (ModuleNotFoundError, AttributeError, ValueError, ImportError) as e:
         logger.warning(e)
         continue
 

mlip_arena/models/externals/fairchem.py

@@ -10,6 +10,33 @@ 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 eSEN(OCPCalculator):
+    def __init__(
+        self,
+        checkpoint=REGISTRY["eSEN"]["checkpoint"],
+        cache_dir=None,
+        cpu=False, # TODO: cannot assign device
+        seed=0,
+        **kwargs,
+    ) -> None:
+
+        # https://huggingface.co/facebook/OMAT24/resolve/main/esen_30m_oam.pt
+
+        checkpoint_path = hf_hub_download(
+            "fairchem/OMAT24",
+            filename=checkpoint,
+            revision="13ab5b8d71af67bd1c83fbbf53250c82cd87f506",
+            cache_dir=cache_dir
+        )
+        kwargs.pop("device", None)
+        super().__init__(
+            checkpoint_path=checkpoint_path,
+            cpu=cpu,
+            seed=seed,
+            **kwargs,
+        )
+
 class eqV2(OCPCalculator):
     def __init__(
         self,

mlip_arena/models/externals/matgl.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import matgl
 import torch
 from matgl.ext.ase import PESCalculator
+from typing import Literal
 
 
 class M3GNet(PESCalculator):
@@ -11,8 +12,10 @@ class M3GNet(PESCalculator):
         checkpoint="M3GNet-MP-2021.2.8-PES",
         # TODO: cannot assign device
         state_attr: torch.Tensor | None = None,
+        stress_unit: Literal["eV/A3", "GPa"] = "GPa",
         stress_weight: float = 1.0,
+        use_voigt: bool = False,
         **kwargs,
     ) -> None:
         potential = matgl.load_model(checkpoint)
-        super().__init__(potential, state_attr, stress_weight, **kwargs)
+        super().__init__(potential, state_attr, stress_unit, stress_weight, use_voigt, **kwargs)

mlip_arena/models/registry.yaml

@@ -192,6 +192,29 @@ MACE-MPA:
   npt: true
   license: MIT
 
+eSEN:
+  module: externals
+  class: eSEN
+  family: fairchem
+  package: fairchem-core==1.9.0
+  checkpoint: esen_30m_oam.pt
+  username: fairchem # HF handle
+  last-update: 2025-04-21
+  datetime: 2025-04-21
+  datasets:
+    - OMat
+    - MPTrj
+    - Alexandria
+  gpu-tasks:
+    - homonuclear-diatomics
+  prediction: EFS
+  nvt: true
+  npt: false # https://github.com/FAIR-Chem/fairchem/issues/888, https://github.com/atomind-ai/mlip-arena/issues/17
+  date: 2025-04-14
+  github: https://github.com/FAIR-Chem/fairchem
+  doi: https://arxiv.org/abs/2502.12147
+  license: Modified Apache-2.0 (Meta)
+
 EquiformerV2(OC22):
   module: externals
   class: EquiformerV2

mlip_arena/tasks/diatomics/analysis.py

@@ -0,0 +1,198 @@
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+from ase.data import chemical_symbols
+from ase.io import read
+from scipy import stats
+from scipy.interpolate import UnivariateSpline
+from tqdm.auto import tqdm
+
+from mlip_arena.models import REGISTRY, MLIPEnum
+
+for model in MLIPEnum:
+
+    df = pd.DataFrame(
+        columns=[
+            "name",
+            "method",
+            "R",
+            "E",
+            "F",
+            "S^2",
+            "force-flip-times",
+            "force-total-variation",
+            "force-jump",
+            "energy-diff-flip-times",
+            "energy-grad-norm-max",
+            "energy-jump",
+            "energy-total-variation",
+            "tortuosity",
+            "conservation-deviation",
+            "spearman-descending-force",
+            "spearman-ascending-force",
+            "spearman-repulsion-energy",
+            "spearman-attraction-energy",
+            "pbe-energy-mae",
+            "pbe-force-mae",
+        ]
+    )
+
+    for symbol in tqdm(chemical_symbols[1:]):
+        da = symbol + symbol
+
+        out_dir = Path(model.name)
+
+        traj_fpath = out_dir / f"{str(da)}.extxyz"
+
+        if traj_fpath.exists():
+            traj = read(traj_fpath, index=":")
+        else:
+            continue
+
+        Rs, Es, Fs, S2s = [], [], [], []
+        for atoms in traj:
+            vec = atoms.positions[1] - atoms.positions[0]
+            r = np.linalg.norm(vec)
+            e = atoms.get_potential_energy()
+            f = np.inner(vec / r, atoms.get_forces()[1])
+            # s2 = np.mean(np.power(atoms.get_magnetic_moments(), 2))
+
+            Rs.append(r)
+            Es.append(e)
+            Fs.append(f)
+            # S2s.append(s2)
+
+        rs = np.array(Rs)
+        es = np.array(Es)
+        fs = np.array(Fs)
+
+        # sort interatomic distances and align to zero at far field
+        indices = np.argsort(rs)[::-1]
+        rs = rs[indices]
+        es = es[indices]
+        eshift = es[0]
+        es -= eshift
+        fs = fs[indices]
+
+        iminf = np.argmin(fs)
+        imine = np.argmin(es)
+
+        de_dr = np.gradient(es, rs)
+        d2e_dr2 = np.gradient(de_dr, rs)
+
+        # avoid numerical sensitity close to zero
+        rounded_fs = np.copy(fs)
+        rounded_fs[np.abs(rounded_fs) < 1e-2] = 0  # 10meV/A
+        fs_sign = np.sign(rounded_fs)
+        mask = fs_sign != 0
+        rounded_fs = rounded_fs[mask]
+        fs_sign = fs_sign[mask]
+        f_flip = np.diff(fs_sign) != 0
+
+        fdiff = np.diff(fs)
+        fdiff_sign = np.sign(fdiff)
+        mask = fdiff_sign != 0
+        fdiff = fdiff[mask]
+        fdiff_sign = fdiff_sign[mask]
+        fdiff_flip = np.diff(fdiff_sign) != 0
+        fjump = (
+            np.abs(fdiff[:-1][fdiff_flip]).sum() + np.abs(fdiff[1:][fdiff_flip]).sum()
+        )
+
+        ediff = np.diff(es)
+        ediff[np.abs(ediff) < 1e-3] = 0  # 1meV
+        ediff_sign = np.sign(ediff)
+        mask = ediff_sign != 0
+        ediff = ediff[mask]
+        ediff_sign = ediff_sign[mask]
+        ediff_flip = np.diff(ediff_sign) != 0
+        ejump = (
+            np.abs(ediff[:-1][ediff_flip]).sum() + np.abs(ediff[1:][ediff_flip]).sum()
+        )
+
+        try:
+            pbe_traj = read(f"./vasp/{da}/PBE.extxyz", index=":")
+
+            pbe_rs, pbe_es, pbe_fs = [], [], []
+
+            for atoms in pbe_traj:
+                vec = atoms.positions[1] - atoms.positions[0]
+                r = np.linalg.norm(vec)
+                pbe_rs.append(r)
+                pbe_es.append(atoms.get_potential_energy())
+                pbe_fs.append(np.inner(vec / r, atoms.get_forces()[1]))
+
+            pbe_rs = np.array(pbe_rs)
+            pbe_es = np.array(pbe_es)
+            pbe_fs = np.array(pbe_fs)
+
+            indices = np.argsort(pbe_rs)
+            pbe_rs = pbe_rs[indices]
+            pbe_es = pbe_es[indices]
+            pbe_fs = pbe_fs[indices]
+
+            pbe_es -= pbe_es[-1]
+
+            xs = np.linspace(pbe_rs.min(), pbe_rs.max(), int(1e3))
+
+            cs = UnivariateSpline(pbe_rs, pbe_es, s=0)
+            pbe_energy_mae = np.mean(np.abs(es - cs(rs)))
+
+            cs = UnivariateSpline(pbe_rs, pbe_fs, s=0)
+            pbe_force_mae = np.mean(np.abs(fs - cs(rs)))
+        except Exception as e:
+            print(e)
+            pbe_energy_mae = None
+            pbe_force_mae = None
+
+        conservation_deviation = np.mean(np.abs(fs + de_dr))
+
+        etv = np.sum(np.abs(np.diff(es)))
+
+        data = {
+            "name": da,
+            "method": model.name,
+            "R": rs,
+            "E": es + eshift,
+            "F": fs,
+            "S^2": S2s,
+            "force-flip-times": np.sum(f_flip),
+            "force-total-variation": np.sum(np.abs(np.diff(fs))),
+            "force-jump": fjump,
+            "energy-diff-flip-times": np.sum(ediff_flip),
+            "energy-grad-norm-max": np.max(np.abs(de_dr)),
+            "energy-jump": ejump,
+            # "energy-grad-norm-mean": np.mean(de_dr_abs),
+            "energy-total-variation": etv,
+            "tortuosity": etv / (abs(es[0] - es.min()) + (es[-1] - es.min())),
+            "conservation-deviation": conservation_deviation,
+            "spearman-descending-force": stats.spearmanr(
+                rs[iminf:], fs[iminf:]
+            ).statistic,
+            "spearman-ascending-force": stats.spearmanr(
+                rs[:iminf], fs[:iminf]
+            ).statistic,
+            "spearman-repulsion-energy": stats.spearmanr(
+                rs[imine:], es[imine:]
+            ).statistic,
+            "spearman-attraction-energy": stats.spearmanr(
+                rs[:imine], es[:imine]
+            ).statistic,
+            "pbe-energy-mae": pbe_energy_mae,
+            "pbe-force-mae": pbe_force_mae,
+        }
+
+        df = pd.concat([df, pd.DataFrame([data])], ignore_index=True)
+
+    json_fpath = Path(REGISTRY[model.name]["family"]) / "homonuclear-diatomics.json"
+
+    if json_fpath.exists():
+        df0 = pd.read_json(json_fpath)
+        df = pd.concat([df0, df], ignore_index=True)
+        df.drop_duplicates(inplace=True, subset=["name", "method"], keep="last")
+
+    df.to_json(json_fpath, orient="records")
+
+    json_fpath = Path(model.name) / "homonuclear-diatomics.json"
+    df.to_json(json_fpath, orient="records")

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

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:06fe63898d03644e491ef12e275884c110cf993b8182403408d65fdd23e51816
-size 2136895
+oid sha256:5f955a5e08dbb9be3d0a04f06bd4cf93ac63a6a401132a7e74e971365133472a
+size 4286019

mlip_arena/tasks/diatomics/run.py

@@ -0,0 +1,131 @@
+import itertools
+from pathlib import Path
+
+import numpy as np
+from ase import Atom, Atoms
+from ase.calculators.calculator import BaseCalculator
+from ase.data import chemical_symbols, covalent_radii, vdw_alvarez
+from ase.io import read, write
+from prefect import flow, task
+from tqdm.auto import tqdm
+
+from mlip_arena.models import REGISTRY, MLIPEnum
+from mlip_arena.tasks.utils import get_calculator
+
+
+@task
+def homonuclear_diatomics(symbol: str, calculator: BaseCalculator, out_dir: Path):
+    """
+    Calculate potential energy curves for homonuclear diatomic molecules.
+
+    This function computes the potential energy of a diatomic molecule (two atoms of
+    the same element) across a range of interatomic distances. The distance range is
+    automatically determined from the covalent and van der Waals radii of the element.
+
+    Args:
+        symbol: Chemical symbol of the atom (e.g., 'H', 'O', 'Fe')
+        calculator: ASE calculator object used to compute the potential energies. Could be VASP, MLIP, etc.
+
+    Returns:
+        None: Results are saved as trajectory files in a directory structure:
+              /{model_family}/{element_pair}/{model_name}.extxyz
+
+    Note:
+        - Minimum distance is set to 0.9× the covalent radius
+        - Maximum distance is set to 3.1× the van der Waals radius (or 6 Å if unknown)
+        - Distance step size is fixed at 0.01 Å
+        - If an existing trajectory file is found, the calculation will resume from where it left off
+        - The atoms are placed in a periodic box large enough to avoid self-interaction
+    """
+
+    atom = Atom(symbol)
+    rmin = 0.9 * covalent_radii[atom.number]
+    rvdw = (
+        vdw_alvarez.vdw_radii[atom.number]
+        if atom.number < len(vdw_alvarez.vdw_radii)
+        else np.nan
+    )
+    rmax = 3.1 * rvdw if not np.isnan(rvdw) else 6
+    rstep = 0.01
+    npts = int((rmax - rmin) / rstep)
+
+    rs = np.linspace(rmin, rmax, npts)
+    es = np.zeros_like(rs)
+
+    da = symbol + symbol
+
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    skip = 0
+
+    a = 5 * rmax
+    r = rs[0]
+
+    positions = [
+        [a / 2 - r / 2, a / 2, a / 2],
+        [a / 2 + r / 2, a / 2, a / 2],
+    ]
+
+    traj_fpath = out_dir / f"{da!s}.extxyz"
+
+    if traj_fpath.exists():
+        traj = read(traj_fpath, index=":")
+        skip = len(traj)
+        atoms = traj[-1]
+    else:
+        # Create the unit cell with two atoms
+        atoms = Atoms(
+            da,
+            positions=positions,
+            # magmoms=magmoms,
+            cell=[a, a + 0.001, a + 0.002],
+            pbc=False,
+        )
+
+    atoms.calc = calculator
+
+    for i, r in enumerate(tqdm(rs)):
+        if i < skip:
+            continue
+
+        positions = [
+            [a / 2 - r / 2, a / 2, a / 2],
+            [a / 2 + r / 2, a / 2, a / 2],
+        ]
+
+        # atoms.set_initial_magnetic_moments(magmoms)
+        atoms.set_positions(positions)
+        es[i] = atoms.get_potential_energy()
+        write(traj_fpath, atoms, append="a")
+
+
+@flow
+def submit_homonuclear_diatomics():
+    futures = []
+    for symbol, model in itertools.product(
+        chemical_symbols[1:],
+        MLIPEnum,
+    ):
+        if "homonuclear-diatomics" not in REGISTRY[model.name].get("gpu-tasks", []):
+            continue
+
+        out_dir = Path(__file__).parent / model.name
+
+        calculator = get_calculator(model)
+
+        # if not (out_dir / "homonuclear-diatomics.json").exists():
+        future = homonuclear_diatomics.submit(
+            symbol,
+            calculator,
+            out_dir=out_dir,
+        )
+        futures.append(future)
+
+    return [f.result(raise_on_failure=False) for f in futures]
+
+
+if __name__ == "__main__":
+    submit_homonuclear_diatomics.with_options(
+        # task_runner=DaskTaskRunner(address=client.scheduler.address),
+        log_prints=True,
+    )()

pyproject.toml

@@ -47,19 +47,14 @@ app = [
     "plotly",
 ]
 test = [
-    "torch==2.2.0",
+    "torch==2.4.0",
     "torch_dftd==0.4.0",
-    "nvidia-ml-py==12.560.30",
     "e3nn==0.5.0",
-    "matgl==1.1.2",
-    # "dgl==2.4.0",
+    "dgl",
     "chgnet==0.3.8",
-    "fairchem-core==1.2.0",
     "sevenn==0.9.3.post1",
-    "orb-models==0.4.0",
-    "pynanoflann@git+https://github.com/dwastberg/pynanoflann#egg=af434039ae14bedcbb838a7808924d6689274168",
     "alignn==2024.5.27",
-    "mattersim==1.0.0rc9",
+    "mattersim==1.1.2",
     "torchani==2.2.4",
     "pytest",
     "pytest-xdist",
@@ -69,7 +64,18 @@ test = [
     "streamlit==1.43.2"
 ]
 mace = [
-    "mace-torch==0.3.9",
+    "mace-torch==0.3.12",
+]
+matgl = [
+    "matgl==1.2.6",
+]
+fairchem = [
+    "hydra-core",
+    "fairchem-core@git+https://github.com/facebookresearch/fairchem.git#subdirectory=packages/fairchem-core",
+]
+orb = [
+    "orb-models==0.4.0",
+    "pynanoflann@git+https://github.com/dwastberg/pynanoflann#egg=af434039ae14bedcbb838a7808924d6689274168",
 ]
 deepmd = [
     "torch==2.2.0",

scripts/install-linux.sh

@@ -1,8 +1,10 @@
-TORCH=2.2
-CUDA=cu121
+TORCH=2.4
+CUDA=cu124
 uv pip install torch==${TORCH}.0
-uv pip install torch-scatter -f https://data.pyg.org/whl/torch-${TORCH}.0+${CUDA}.html
-uv pip install torch-sparse -f https://data.pyg.org/whl/torch-${TORCH}.0+${CUDA}.html
+uv pip install torch-scatter torch-sparse -f https://data.pyg.org/whl/torch-${TORCH}.0+${CUDA}.html
 uv pip install dgl -f https://data.dgl.ai/wheels/torch-${TORCH}/${CUDA}/repo.html
+uv pip install -e .[fairchem]
+uv pip install -e .[orb]
+uv pip install -e .[matgl]
 uv pip install -e .[test]
 uv pip install -e .[mace]

serve/tasks/homonuclear-diatomics.py

@@ -5,10 +5,10 @@ import pandas as pd
 import plotly.colors as pcolors
 import plotly.graph_objects as go
 import streamlit as st
-from mlip_arena.models import REGISTRY
+from ase.data import chemical_symbols
 from plotly.subplots import make_subplots
 
-from ase.data import chemical_symbols
+from mlip_arena.models import REGISTRY
 
 st.markdown(
     """
@@ -29,16 +29,7 @@ valid_models = [
 mlip_methods = container.multiselect(
     "MLIPs",
     valid_models,
-    [
-        "MACE-MP(M)",
-        "CHGNet",
-        "M3GNet",
-        "MatterSim",
-        "SevenNet",
-        "ORBv2",
-        "eqV2(OMat)",
-        "ANI2x",
-    ],
+    ["MACE-MP(M)", "MatterSim", "SevenNet", "ORBv2", "eqV2(OMat)", "ANI2x", "eSEN"],
 )
 dft_methods = container.multiselect("DFT Methods", ["PBE"], ["PBE"])