Add neb task (#38)

* add neb task

* return neb tool fit object

* add run from endpoints

* add neb test

* fix typo

* try parallel pytest

* add pytest-xdist

.github/workflows/test.yaml

@@ -26,7 +26,7 @@ jobs:
     - name: Install uv
       uses: astral-sh/setup-uv@v4
       with:
-        enable-cahce: true
+        enable-cache: true
         cache-dependency-glob: "pyproject.toml"
 
     - name: Set up Python ${{ matrix.python-version }}
@@ -63,4 +63,4 @@ jobs:
         PREFECT_API_KEY: ${{ secrets.PREFECT_API_KEY }}
         PREFECT_API_URL: ${{ secrets.PREFECT_API_URL }}
       run: |
-        pytest -vvra tests
+        pytest -vra tests -n 5

mlip_arena/tasks/neb.py

@@ -0,0 +1,268 @@
+"""
+Defines nudged elastic band (NEB) task
+
+This module has been modified from MatCalc
+https://github.com/materialsvirtuallab/matcalc/blob/main/src/matcalc/neb.py
+
+https://github.com/materialsvirtuallab/matcalc/blob/main/LICENSE
+
+BSD 3-Clause License
+
+Copyright (c) 2023, Materials Virtual Lab
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Literal
+
+from prefect import task
+from prefect.cache_policies import INPUTS, TASK_SOURCE
+from prefect.runtime import task_run
+from prefect.states import State
+
+from ase import Atoms
+from ase.filters import *  # type: ignore
+from ase.mep.neb import NEB, NEBTools
+from ase.optimize import *  # type: ignore
+from ase.optimize.optimize import Optimizer
+from ase.utils.forcecurve import fit_images
+from mlip_arena.models import MLIPEnum
+from mlip_arena.tasks.optimize import run as OPT
+from mlip_arena.tasks.utils import get_calculator
+from pymatgen.io.ase import AseAtomsAdaptor
+
+if TYPE_CHECKING:
+    pass
+
+_valid_optimizers: dict[str, Optimizer] = {
+    "MDMin": MDMin,
+    "FIRE": FIRE,
+    "FIRE2": FIRE2,
+    "LBFGS": LBFGS,
+    "LBFGSLineSearch": LBFGSLineSearch,
+    "BFGS": BFGS,
+    # "BFGSLineSearch": BFGSLineSearch, # NEB does not support BFGSLineSearch
+    "QuasiNewton": QuasiNewton,
+    "GPMin": GPMin,
+    "CellAwareBFGS": CellAwareBFGS,
+    "ODE12r": ODE12r,
+}  # type: ignore
+
+
+def _generate_task_run_name():
+    task_name = task_run.task_name
+    parameters = task_run.parameters
+
+    if "images" in parameters:
+        atoms = parameters["images"][0]
+    elif "start" in parameters:
+        atoms = parameters["start"]
+    else:
+        raise ValueError("No images or start atoms found in parameters")
+    
+    calculator_name = parameters["calculator_name"]
+
+    return f"{task_name}: {atoms.get_chemical_formula()} - {calculator_name}"
+
+
+@task(
+    name="NEB from images",
+    task_run_name=_generate_task_run_name,
+    cache_policy=TASK_SOURCE + INPUTS,
+)
+def run(
+    images: list[Atoms],
+    calculator_name: str | MLIPEnum,
+    calculator_kwargs: dict | None = None,
+    dispersion: str | None = None,
+    dispersion_kwargs: dict | None = None,
+    device: str | None = None,
+    optimizer: Optimizer | str = "MDMin",  # type: ignore
+    optimizer_kwargs: dict | None = None,
+    criterion: dict | None = None,
+    interpolation: Literal["linear", "idpp"] = "idpp",
+    climb: bool = True,
+    traj_file: str | Path | None = None,
+) -> dict[str, Any] | State:
+    """Run the nudged elastic band (NEB) calculation.
+
+    Args:
+        images (list[Atoms]): The images.
+        calculator_name (str | MLIPEnum): The calculator name.
+        calculator_kwargs (dict, optional): The calculator kwargs. Defaults to None.
+        dispersion (str, optional): The dispersion. Defaults to None.
+        dispersion_kwargs (dict, optional): The dispersion kwargs. Defaults to None.
+        device (str, optional): The device. Defaults to None.
+        optimizer (Optimizer | str, optional): The optimizer. Defaults to "BFGSLineSearch".
+        optimizer_kwargs (dict, optional): The optimizer kwargs. Defaults to None.
+        criterion (dict, optional): The criterion. Defaults to None.
+        interpolation (Literal['linear', 'idpp'], optional): The interpolation method. Defaults to "idpp".
+        climb (bool, optional): Whether to use the climbing image. Defaults to True.
+        traj_file (str | Path, optional): The trajectory file. Defaults to None.
+
+    Returns:
+        dict[str, Any] | State: The energy barrier.
+    """
+
+    calc = get_calculator(
+        calculator_name,
+        calculator_kwargs,
+        dispersion=dispersion,
+        dispersion_kwargs=dispersion_kwargs,
+        device=device,
+    )
+
+    for image in images:
+        assert isinstance(image, Atoms)
+        image.calc = calc
+
+    neb = NEB(images, climb=climb, allow_shared_calculator=True)
+
+    neb.interpolate(method=interpolation)
+
+    if isinstance(optimizer, str):
+        if optimizer not in _valid_optimizers:
+            raise ValueError(f"Invalid optimizer: {optimizer}")
+        optimizer = _valid_optimizers[optimizer]
+
+    optimizer_kwargs = optimizer_kwargs or {}
+    criterion = criterion or {}
+
+    optimizer_instance = optimizer(neb, trajectory=traj_file, **optimizer_kwargs)  # type: ignore
+
+    optimizer_instance.run(**criterion)
+
+    neb_tool = NEBTools(neb.images)
+
+    return {
+        "barrier": neb_tool.get_barrier(),
+        "images": neb.images,
+        "forcefit": fit_images(neb.images),
+    }
+
+
+@task(
+    name="NEB from end points",
+    task_run_name=_generate_task_run_name,
+    cache_policy=TASK_SOURCE + INPUTS,
+)
+def run_from_end_points(
+    start: Atoms,
+    end: Atoms,
+    n_images: int,
+    calculator_name: str | MLIPEnum,
+    calculator_kwargs: dict | None = None,
+    dispersion: str | None = None,
+    dispersion_kwargs: dict | None = None,
+    device: str | None = None,
+    optimizer: Optimizer | str = "BFGS",  # type: ignore
+    optimizer_kwargs: dict | None = None,
+    criterion: dict | None = None,
+    relax_end_points: bool = True,
+    interpolation: Literal["linear", "idpp"] = "idpp",
+    climb: bool = True,
+    traj_file: str | Path | None = None,
+) -> dict[str, Any] | State:
+    """Run the nudged elastic band (NEB) calculation from end points.
+
+    Args:
+        start (Atoms): The start image.
+        end (Atoms): The end image.
+        n_images (int): The number of images.
+        calculator_name (str | MLIPEnum): The calculator name.
+        calculator_kwargs (dict, optional): The calculator kwargs. Defaults to None.
+        dispersion (str, optional): The dispersion. Defaults to None.
+        dispersion_kwargs (dict, optional): The dispersion kwargs. Defaults to None.
+        device (str, optional): The device. Defaults to None.
+        optimizer (Optimizer | str, optional): The optimizer. Defaults to "BFGSLineSearch".
+        optimizer_kwargs (dict, optional): The optimizer kwargs. Defaults to None.
+        criterion (dict, optional): The criterion. Defaults to None.
+        interpolation (Literal['linear', 'idpp'], optional): The interpolation method. Defaults to "idpp".
+        climb (bool, optional): Whether to use the climbing image. Defaults to True.
+        traj_file (str | Path, optional): The trajectory file. Defaults to None.
+
+    Returns:
+        dict[str, Any] | State: The energy barrier.
+    """
+
+    if relax_end_points:
+        relax = OPT(
+            atoms=start.copy(),
+            calculator_name=calculator_name,
+            calculator_kwargs=calculator_kwargs,
+            dispersion=dispersion,
+            dispersion_kwargs=dispersion_kwargs,
+            device=device,
+            optimizer=optimizer,
+            optimizer_kwargs=optimizer_kwargs,
+            criterion=criterion,
+        )
+        start = relax["atoms"]
+
+        relax = OPT(
+            atoms=end.copy(),
+            calculator_name=calculator_name,
+            calculator_kwargs=calculator_kwargs,
+            dispersion=dispersion,
+            dispersion_kwargs=dispersion_kwargs,
+            device=device,
+            optimizer=optimizer,
+            optimizer_kwargs=optimizer_kwargs,
+            criterion=criterion,
+        )
+        end = relax["atoms"]
+
+    path = (
+        AseAtomsAdaptor()
+        .get_structure(start)
+        .interpolate(
+            AseAtomsAdaptor().get_structure(end),
+            nimages=n_images - 1,
+            interpolate_lattices=False,
+            pbc=False,
+            autosort_tol=0.5,
+        )
+    )
+
+    images = [s.to_ase_atoms() for s in path]
+
+    return run(
+        images,
+        calculator_name,
+        calculator_kwargs=calculator_kwargs,
+        dispersion=dispersion,
+        dispersion_kwargs=dispersion_kwargs,
+        device=device,
+        optimizer=optimizer,
+        optimizer_kwargs=optimizer_kwargs,
+        criterion=criterion,
+        interpolation=interpolation,
+        climb=climb,
+        traj_file=traj_file,
+    )

mlip_arena/tasks/utils.py

@@ -0,0 +1,51 @@
+"""Utility functions for MLIP models."""
+
+from __future__ import annotations
+
+from torch_dftd.torch_dftd3_calculator import TorchDFTD3Calculator
+
+from ase.calculators.calculator import Calculator
+from ase.calculators.mixing import SumCalculator
+from ase.filters import *  # type: ignore
+from ase.optimize import *  # type: ignore
+from mlip_arena.models import MLIPEnum
+from mlip_arena.models.utils import get_freer_device
+
+
+def get_calculator(
+    calculator_name: str | MLIPEnum,
+    calculator_kwargs: dict | None,
+    dispersion: str | None = None,
+    dispersion_kwargs: dict | None = None,
+    device: str | None = None,
+) -> Calculator:
+    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 hasattr(MLIPEnum, calculator_name):
+        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}")
+
+    assert isinstance(calc, Calculator)
+    return calc

pyproject.toml

@@ -64,6 +64,7 @@ test = [
     "alignn==2024.5.27",
     "mattersim==1.0.0rc9",
     "pytest",
+    "pytest-xdist",
     "prefect>=3.0.4,<3.1.0",
 ]
 mace = [

tests/test_neb.py

@@ -0,0 +1,45 @@
+import sys
+
+import pytest
+from mlip_arena.models import MLIPEnum
+from mlip_arena.tasks.neb import run as NEB
+from mlip_arena.tasks.neb import run_from_end_points as NEB
+from prefect.testing.utilities import prefect_test_harness
+
+from ase.spacegroup import crystal
+
+pristine = crystal(
+    "Al", [(0, 0, 0)], spacegroup=225, cellpar=[4.05, 4.05, 4.05, 90, 90, 90]
+) * (3, 3, 3)
+
+atoms = pristine.copy()
+del atoms[0]
+start = atoms.copy()
+
+atoms = pristine.copy()
+del atoms[1]
+end = atoms.copy()
+
+
+@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_neb(model: MLIPEnum):
+    """
+    Test NEB prefect workflow with a simple cubic lattice.
+    """
+
+    with prefect_test_harness():
+        result = NEB(
+            start=start.copy(),
+            end=end.copy(),
+            n_images=5,
+            calculator_name=model.name,
+            optimizer="FIRE2",
+        )
+
+        assert isinstance(result, dict)
+        assert isinstance(result["barrier"][0], float)
+        assert isinstance(result["barrier"][1], float)