Nudged Elastic BandΒΆ

mlipx provides a command line interface to interpolate and create a NEB path from inital-final or initial-ts-final images and run NEB on the interpolated images. You can run the following command to instantiate a test directory:

(.venv) $ mlipx recipes neb --models mace-mpa-0,7net-0,7net-mf-ompa-mpa,mattersim,grace-2l-omat,chgnet --datapath ../../data/neb_end_p.xyz --repro
(.venv) $ mlipx compare --glob "*NEBs"

This test uses the following Nodes together with your provided model in the models.py file:

Content of main.py
import zntrack
from models import MODELS

import mlipx

DATAPATH = "../../data/neb_end_p.xyz"

project = zntrack.Project()

with project.group("initialize"):
    data = mlipx.LoadDataFile(path=DATAPATH)
    trajectory = mlipx.NEBinterpolate(data=data.frames, n_images=5, mic=True)

for model_name, model in MODELS.items():
    with project.group(model_name):
        neb = mlipx.NEBs(
            data=trajectory.frames,
            model=model,
            relax=True,
            fmax=0.05,
        )

project.build()
Content of models.py
import dataclasses

import mlipx
from mlipx.nodes.generic_ase import Device

ALL_MODELS = {}

# https://github.com/ACEsuit/mace
ALL_MODELS["mace-mpa-0"] = mlipx.GenericASECalculator(
    module="mace.calculators",
    class_name="mace_mp",
    device="auto",
    kwargs={"model": "../../models/mace-mpa-0-medium.model"}
    # MLIPX-hub model path, adjust as needed
)
# https://github.com/MDIL-SNU/SevenNet
ALL_MODELS["7net-0"] = mlipx.GenericASECalculator(
    module="sevenn.sevennet_calculator",
    class_name="SevenNetCalculator",
    device="auto",
    kwargs={"model": "7net-0"}
)
ALL_MODELS["7net-mf-ompa-mpa"] = mlipx.GenericASECalculator(
    module="sevenn.sevennet_calculator",
    class_name="SevenNetCalculator",
    device="auto",
    kwargs={"model": "7net-mf-ompa", "modal": "mpa"}
)

# https://github.com/orbital-materials/orb-models
@dataclasses.dataclass
class OrbCalc:
    name: str
    device: Device | None = None
    kwargs: dict = dataclasses.field(default_factory=dict)

    def get_calculator(self, **kwargs):
        from orb_models.forcefield import pretrained
        from orb_models.forcefield.calculator import ORBCalculator

        method = getattr(pretrained, self.name)
        if self.device is None:
            orbff = method(**self.kwargs)
            calc = ORBCalculator(orbff, **self.kwargs)
        elif self.device == Device.AUTO:
            orbff = method(device=Device.resolve_auto(), **self.kwargs)
            calc = ORBCalculator(orbff, device=Device.resolve_auto(), **self.kwargs)
        else:
            orbff = method(device=self.device, **self.kwargs)
            calc = ORBCalculator(orbff, device=self.device, **self.kwargs)
        return calc

    @property
    def available(self) -> bool:
        try:
            from orb_models.forcefield import pretrained
            from orb_models.forcefield.calculator import ORBCalculator
            return True
        except ImportError:
            return False

ALL_MODELS["orb-v2"] = OrbCalc(
    name="orb_v2",
    device="auto"
)
ALL_MODELS["orb-v3"] = OrbCalc(
    name="orb_v3_conservative_inf_omat",
    device="auto"
)

# https://github.com/CederGroupHub/chgnet
ALL_MODELS["chgnet"] = mlipx.GenericASECalculator(
    module="chgnet.model",
    class_name="CHGNetCalculator",
)
# https://github.com/microsoft/mattersim
ALL_MODELS["mattersim"] = mlipx.GenericASECalculator(
    module="mattersim.forcefield",
    class_name="MatterSimCalculator",
    device="auto",
)
# https://www.faccts.de/orca/
ALL_MODELS["orca"] = mlipx.OrcaSinglePoint(
    orcasimpleinput= "PBE def2-TZVP TightSCF EnGrad",
    orcablocks ="%pal nprocs 8 end",
    orca_shell="",
)

# https://gracemaker.readthedocs.io/en/latest/gracemaker/foundation/
ALL_MODELS["grace-2l-omat"] = mlipx.GenericASECalculator(
    module="tensorpotential.calculator",
    class_name="TPCalculator",
    device=None,
    kwargs={
        "model": "../../models/GRACE-2L-OMAT",
    },
    # MLIPX-hub model path, adjust as needed
)

# OPTIONAL
# ========
# If you have custom property names you can use the UpdatedFramesCalc
# to set the energy, force and isolated_energies keys mlipx expects.

# REFERENCE = mlipx.UpdateFramesCalc(
#     results_mapping={"energy": "DFT_ENERGY", "forces": "DFT_FORCES"},
#     info_mapping={mlipx.abc.ASEKeys.isolated_energies.value: "isol_ene"},
# )

# ============================================================
# THE SELECTED MODELS!
# ONLY THESE MODELS WILL BE USED IN THE RECIPE
# ============================================================
MODELS = {
    "mace-mpa-0": ALL_MODELS["mace-mpa-0"],
    "7net-0": ALL_MODELS["7net-0"],
    "7net-mf-ompa-mpa": ALL_MODELS["7net-mf-ompa-mpa"],
    "mattersim": ALL_MODELS["mattersim"],
    "grace-2l-omat": ALL_MODELS["grace-2l-omat"],
    "chgnet": ALL_MODELS["chgnet"],
}