Enforce copying atoms and refactor calculator instantiation to allow custom calculator (#47)

* enforce copying atoms; refactor calculator parsings

* refactor test

* fix `generate_task_run_name`

* update readme example

* loosen a tiny bit pytest approx

* update md example to apply dispersion correction

.github/README.md

@@ -12,9 +12,9 @@
 > [!NOTE]
 > Contributions of new tasks are very welcome! If you're interested in joining the effort, please reach out to Yuan at [[email protected]](mailto:[email protected]). See [project page](https://github.com/orgs/atomind-ai/projects/1) for some outstanding tasks, or propose new one in [Discussion](https://github.com/atomind-ai/mlip-arena/discussions/new?category=ideas).
 
-MLIP Arena is a unified platform for evaluating foundation machine learning interatomic potentials (MLIPs) beyond conventional error metrics. It focuses on revealing the underlying physics and chemistry learned by these models and assessing their utilitarian performance agnostic to underlying model architecture. The platform's benchmarks are specifically designed to evaluate the readiness and reliability of open-source, open-weight models in accurately reproducing both qualitative and quantitative behaviors of atomic systems.
+MLIP Arena is a unified platform for evaluating foundation machine learning interatomic potentials (MLIPs) beyond conventional error metrics. It focuses on revealing the physics and chemistry learned by these models and assessing their utilitarian performance agnostic to underlying model architecture. The platform's benchmarks are specifically designed to evaluate the readiness and reliability of open-source, open-weight models in accurately reproducing both qualitative and quantitative behaviors of atomic systems.
 
-MLIP Arena leverages modern pythonic workflow orchestractor [Prefect](https://www.prefect.io/) to enable advanced task/flow chaining and caching.
+MLIP Arena leverages modern pythonic workflow orchestrator [Prefect](https://www.prefect.io/) to enable advanced task/flow chaining and caching.
 
 ## Installation
 
@@ -46,7 +46,7 @@ DP_ENABLE_TENSORFLOW=0 pip install -e .[deepmd]
 # (Optional) Install uv
 curl -LsSf https://astral.sh/uv/install.sh | sh
 source $HOME/.local/bin/env
-# One script installation
+# One script uv pip installation
 bash scripts/install-macosx.sh
 ```
 
@@ -57,10 +57,12 @@ bash scripts/install-macosx.sh
 Arena provides a unified interface to run all the compiled MLIPs. This can be achieved simply by looping through `MLIPEnum`:
 
 ```python
-from mlip_arena.tasks.md import run as MD 
-# from mlip_arena.tasks import MD # convenient loading
 from mlip_arena.models import MLIPEnum
+from mlip_arena.tasks.md import run as MD 
+# from mlip_arena.tasks import MD # for convenient import
+from mlip_arena.tasks.utils import get_calculator
 
+from ase import units
 from ase.build import bulk
 
 atoms = bulk("Cu", "fcc", a=3.6)
@@ -70,15 +72,18 @@ results = []
 for model in MLIPEnum:
     result = MD(
         atoms=atoms,
-        calculator_name=model,
-        calculator_kwargs={},
+        calculator=get_calculator(
+            model,
+            calculator_kwargs=dict(), # passing into calculator
+            dispersion=True,
+            dispersion_kwargs=dict(damping='bj', xc='pbe', cutoff=40.0 * units.Bohr), # passing into TorchDFTD3Calculator
+        ),
         ensemble="nve",
         dynamics="velocityverlet",
         total_time=1e3, # 1 ps = 1e3 fs
         time_step=2, # fs
     )
     results.append(result)
-
 ```
 
 ## Contribute

mlip_arena/tasks/elasticity.py

@@ -1,7 +1,7 @@
 """
 Defines the tasks for computing the elastic tensor.
 
-This module has been modified from MatCalc 
+This module has been modified from MatCalc
 https://github.com/materialsvirtuallab/matcalc/blob/main/src/matcalc/elasticity.py
 
 https://github.com/materialsvirtuallab/matcalc/blob/main/LICENSE
@@ -41,15 +41,15 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, Any
 
 import numpy as np
+from ase import Atoms
+from ase.calculators.calculator import BaseCalculator
+from ase.optimize.optimize import Optimizer
 from numpy.typing import ArrayLike
 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.optimize.optimize import Optimizer
-from mlip_arena.models import MLIPEnum
 from mlip_arena.tasks.optimize import run as OPT
 from pymatgen.analysis.elasticity import DeformedStructureSet, ElasticTensor, Strain
 from pymatgen.analysis.elasticity.elastic import get_strain_state_dict
@@ -64,7 +64,7 @@ def _generate_task_run_name():
     parameters = task_run.parameters
 
     atoms = parameters["atoms"]
-    calculator_name = parameters["calculator_name"]
+    calculator_name = parameters["calculator"]
 
     return f"{task_name}: {atoms.get_chemical_formula()} - {calculator_name}"
 
@@ -77,11 +77,7 @@ def _generate_task_run_name():
 )
 def run(
     atoms: Atoms,
-    calculator_name: str | MLIPEnum,
-    calculator_kwargs: dict | None = None,
-    dispersion: bool = False,
-    dispersion_kwargs: dict | None = None,
-    device: str | None = None,
+    calculator: BaseCalculator,
     optimizer: Optimizer | str = "BFGSLineSearch",  # type: ignore
     optimizer_kwargs: dict | None = None,
     filter: Filter | str | None = "FrechetCell",  # type: ignore
@@ -97,9 +93,7 @@ def run(
 
     Args:
         atoms (Atoms): The input structure.
-        calculator_name (str | MLIPEnum): The calculator name.
-        calculator_kwargs (dict, optional): The calculator kwargs. Defaults to None.
-        device (str, optional): The device. Defaults to None.
+        calculator (BaseCalculator): The calculator.
         optimizer (Optimizer | str, optional): The optimizer. Defaults to "BFGSLineSearch".
         optimizer_kwargs (dict, optional): The optimizer kwargs. Defaults to None.
         filter (Filter | str, optional): The filter. Defaults to "FrechetCell".
@@ -115,6 +109,8 @@ def run(
         dict[str, Any] | State: The elastic tensor.
     """
 
+    atoms = atoms.copy()
+
     OPT_ = OPT.with_options(
         refresh_cache=not cache_opt,
         persist_result=persist_opt,
@@ -122,11 +118,7 @@ def run(
 
     first_relax = OPT_(
         atoms=atoms,
-        calculator_name=calculator_name,
-        calculator_kwargs=calculator_kwargs,
-        dispersion=dispersion,
-        dispersion_kwargs=dispersion_kwargs,
-        device=device,
+        calculator=calculator,
         optimizer=optimizer,
         optimizer_kwargs=optimizer_kwargs,
         filter=filter,
@@ -172,9 +164,7 @@ def run(
     ]
 
     fit = fit_elastic_tensor(
-        strains,
-        stresses,
-        eq_stress=relaxed.get_stress(voigt=False)
+        strains, stresses, eq_stress=relaxed.get_stress(voigt=False)
     )
 
     return {

mlip_arena/tasks/eos.py

@@ -9,6 +9,9 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, Any
 
 import numpy as np
+from ase import Atoms
+from ase.calculators.calculator import BaseCalculator
+from ase.optimize.optimize import Optimizer
 from prefect import task
 from prefect.cache_policies import INPUTS, TASK_SOURCE
 from prefect.futures import wait
@@ -16,9 +19,6 @@ from prefect.results import ResultRecord
 from prefect.runtime import task_run
 from prefect.states import State
 
-from ase import Atoms
-from ase.optimize.optimize import Optimizer
-from mlip_arena.models import MLIPEnum
 from mlip_arena.tasks.optimize import run as OPT
 from pymatgen.analysis.eos import BirchMurnaghan
 
@@ -31,7 +31,7 @@ def _generate_task_run_name():
     parameters = task_run.parameters
 
     atoms = parameters["atoms"]
-    calculator_name = parameters["calculator_name"]
+    calculator_name = parameters["calculator"]
 
     return f"{task_name}: {atoms.get_chemical_formula()} - {calculator_name}"
 
@@ -41,9 +41,7 @@ def _generate_task_run_name():
 )
 def run(
     atoms: Atoms,
-    calculator_name: str | MLIPEnum,
-    calculator_kwargs: dict | None = None,
-    device: str | None = None,
+    calculator: BaseCalculator,
     optimizer: Optimizer | str = "BFGSLineSearch",  # type: ignore
     optimizer_kwargs: dict | None = None,
     filter: Filter | str | None = "FrechetCell",  # type: ignore
@@ -77,6 +75,8 @@ def run(
         A dictionary containing the EOS data, bulk modulus, equilibrium volume, and equilibrium energy if successful. Otherwise, a prefect state object.
     """
 
+    atoms = atoms.copy()
+
     OPT_ = OPT.with_options(
         refresh_cache=not cache_opt,
         persist_result=cache_opt,
@@ -84,9 +84,7 @@ def run(
 
     state = OPT_(
         atoms=atoms,
-        calculator_name=calculator_name,
-        calculator_kwargs=calculator_kwargs,
-        device=device,
+        calculator=calculator,
         optimizer=optimizer,
         optimizer_kwargs=optimizer_kwargs,
         filter=filter,
@@ -118,9 +116,7 @@ def run(
 
             future = OPT_.submit(
                 atoms=atoms,
-                calculator_name=calculator_name,
-                calculator_kwargs=calculator_kwargs,
-                device=device,
+                calculator=calculator,
                 optimizer=optimizer,
                 optimizer_kwargs=optimizer_kwargs,
                 filter=None,
@@ -144,9 +140,7 @@ def run(
 
             state = OPT_(
                 atoms=atoms,
-                calculator_name=calculator_name,
-                calculator_kwargs=calculator_kwargs,
-                device=device,
+                calculator=calculator,
                 optimizer=optimizer,
                 optimizer_kwargs=optimizer_kwargs,
                 filter=None,
@@ -176,7 +170,6 @@ def run(
 
     return {
         "atoms": relaxed,
-        "calculator_name": calculator_name,
         "eos": {"volumes": volumes, "energies": energies},
         "K": bm.b0_GPa,
         "b0": bm.b0,

mlip_arena/tasks/md.py

@@ -60,14 +60,8 @@ from pathlib import Path
 from typing import Literal
 
 import numpy as np
-from prefect import task
-from prefect.cache_policies import INPUTS, TASK_SOURCE
-from prefect.runtime import task_run
-from scipy.interpolate import interp1d
-from scipy.linalg import schur
-from tqdm.auto import tqdm
-
 from ase import Atoms, units
+from ase.calculators.calculator import BaseCalculator
 from ase.io import read
 from ase.io.trajectory import Trajectory
 from ase.md.andersen import Andersen
@@ -82,8 +76,12 @@ from ase.md.velocitydistribution import (
     ZeroRotation,
 )
 from ase.md.verlet import VelocityVerlet
-from mlip_arena.models import MLIPEnum
-from mlip_arena.tasks.utils import get_calculator
+from prefect import task
+from prefect.cache_policies import INPUTS, TASK_SOURCE
+from prefect.runtime import task_run
+from scipy.interpolate import interp1d
+from scipy.linalg import schur
+from tqdm.auto import tqdm
 
 _valid_dynamics: dict[str, tuple[str, ...]] = {
     "nve": ("velocityverlet",),
@@ -189,7 +187,7 @@ def _generate_task_run_name():
     parameters = task_run.parameters
 
     atoms = parameters["atoms"]
-    calculator_name = parameters["calculator_name"]
+    calculator_name = parameters["calculator"]
 
     return f"{task_name}: {atoms.get_chemical_formula()} - {calculator_name}"
 
@@ -201,11 +199,7 @@ def _generate_task_run_name():
 )
 def run(
     atoms: Atoms,
-    calculator_name: str | MLIPEnum,
-    calculator_kwargs: dict | None = None,
-    dispersion: bool = False,
-    dispersion_kwargs: dict | None = None,
-    device: str | None = None,
+    calculator: BaseCalculator,
     ensemble: Literal["nve", "nvt", "npt"] = "nvt",
     dynamics: str | MolecularDynamics = "langevin",
     time_step: float | None = None,  # fs
@@ -221,13 +215,9 @@ def run(
     restart: bool = True,
 ):
     
-    atoms.calc = get_calculator(
-        calculator_name=calculator_name,
-        calculator_kwargs=calculator_kwargs,
-        dispersion=dispersion,
-        dispersion_kwargs=dispersion_kwargs,
-        device=device,
-    )
+    atoms = atoms.copy()
+
+    atoms.calc = calculator
 
     if time_step is None:
         # If a structure contains an isotope of hydrogen, set default `time_step`

mlip_arena/tasks/neb.py

@@ -41,20 +41,20 @@ from __future__ import annotations
 from pathlib import Path
 from typing import 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.calculators.calculator import BaseCalculator
 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 prefect import task
+from prefect.cache_policies import INPUTS, TASK_SOURCE
+from prefect.runtime import task_run
+from prefect.states import State
+
 from mlip_arena.tasks.optimize import run as OPT
-from mlip_arena.tasks.utils import get_calculator, logger, pformat
+from mlip_arena.tasks.utils import logger, pformat
 from pymatgen.io.ase import AseAtomsAdaptor
 
 _valid_optimizers: dict[str, Optimizer] = {
@@ -83,7 +83,7 @@ def _generate_task_run_name():
     else:
         raise ValueError("No images or start atoms found in parameters")
 
-    calculator_name = parameters["calculator_name"]
+    calculator_name = parameters["calculator"]
 
     return f"{task_name}: {atoms.get_chemical_formula()} - {calculator_name}"
 
@@ -95,11 +95,7 @@ def _generate_task_run_name():
 )
 def run(
     images: list[Atoms],
-    calculator_name: str | MLIPEnum,
-    calculator_kwargs: dict | None = None,
-    dispersion: bool = False,
-    dispersion_kwargs: dict | None = None,
-    device: str | None = None,
+    calculator: BaseCalculator,
     optimizer: Optimizer | str = "MDMin",  # type: ignore
     optimizer_kwargs: dict | None = None,
     criterion: dict | None = None,
@@ -127,17 +123,11 @@ def run(
         dict[str, Any] | State: The energy barrier.
     """
 
-    calc = get_calculator(
-        calculator_name,
-        calculator_kwargs,
-        dispersion=dispersion,
-        dispersion_kwargs=dispersion_kwargs,
-        device=device,
-    )
+    images = [image.copy() for image in images]
 
     for image in images:
         assert isinstance(image, Atoms)
-        image.calc = calc
+        image.calc = calculator
 
     neb = NEB(images, climb=climb, allow_shared_calculator=True)
 
@@ -175,11 +165,7 @@ def run_from_endpoints(
     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,
+    calculator: BaseCalculator,
     optimizer: Optimizer | str = "BFGS",  # type: ignore
     optimizer_kwargs: dict | None = None,
     criterion: dict | None = None,
@@ -216,11 +202,7 @@ def run_from_endpoints(
             refresh_cache=not cache_subtasks,
         )(
             atoms=start.copy(),
-            calculator_name=calculator_name,
-            calculator_kwargs=calculator_kwargs,
-            dispersion=dispersion,
-            dispersion_kwargs=dispersion_kwargs,
-            device=device,
+            calculator=calculator,
             optimizer=optimizer,
             optimizer_kwargs=optimizer_kwargs,
             criterion=criterion,
@@ -231,11 +213,7 @@ def run_from_endpoints(
             refresh_cache=not cache_subtasks,
         )(
             atoms=end.copy(),
-            calculator_name=calculator_name,
-            calculator_kwargs=calculator_kwargs,
-            dispersion=dispersion,
-            dispersion_kwargs=dispersion_kwargs,
-            device=device,
+            calculator=calculator,
             optimizer=optimizer,
             optimizer_kwargs=optimizer_kwargs,
             criterion=criterion,
@@ -260,11 +238,7 @@ def run_from_endpoints(
         refresh_cache=not cache_subtasks,
     )(
         images,
-        calculator_name,
-        calculator_kwargs=calculator_kwargs,
-        dispersion=dispersion,
-        dispersion_kwargs=dispersion_kwargs,
-        device=device,
+        calculator=calculator,
         optimizer=optimizer,
         optimizer_kwargs=optimizer_kwargs,
         criterion=criterion,

mlip_arena/tasks/optimize.py

@@ -4,19 +4,18 @@ Define structure optimization tasks.
 
 from __future__ import annotations
 
-from prefect import task
-from prefect.cache_policies import INPUTS, TASK_SOURCE
-from prefect.runtime import task_run
-
 from ase import Atoms
+from ase.calculators.calculator import BaseCalculator
 from ase.constraints import FixSymmetry
 from ase.filters import *  # type: ignore
 from ase.filters import Filter
 from ase.optimize import *  # type: ignore
 from ase.optimize.optimize import Optimizer
-from mlip_arena.models import MLIPEnum
-from mlip_arena.tasks.utils import get_calculator, logger, pformat
+from prefect import task
+from prefect.cache_policies import INPUTS, TASK_SOURCE
+from prefect.runtime import task_run
 
+from mlip_arena.tasks.utils import logger, pformat
 
 _valid_filters: dict[str, Filter] = {
     "Filter": Filter,
@@ -46,7 +45,7 @@ def _generate_task_run_name():
     parameters = task_run.parameters
 
     atoms = parameters["atoms"]
-    calculator_name = parameters["calculator_name"]
+    calculator_name = parameters["calculator"]
 
     return f"{task_name}: {atoms.get_chemical_formula()} - {calculator_name}"
 
@@ -56,11 +55,7 @@ def _generate_task_run_name():
 )
 def run(
     atoms: Atoms,
-    calculator_name: str | MLIPEnum,
-    calculator_kwargs: dict | None = None,
-    dispersion: bool = False,
-    dispersion_kwargs: dict | None = None,
-    device: str | None = None,
+    calculator: BaseCalculator,
     optimizer: Optimizer | str = BFGSLineSearch,
     optimizer_kwargs: dict | None = None,
     filter: Filter | str | None = None,
@@ -68,13 +63,8 @@ def run(
     criterion: dict | None = None,
     symmetry: bool = False,
 ):
-    atoms.calc = get_calculator(
-        calculator_name=calculator_name,
-        calculator_kwargs=calculator_kwargs,
-        dispersion=dispersion,
-        dispersion_kwargs=dispersion_kwargs,
-        device=device,
-    )
+    atoms = atoms.copy()
+    atoms.calc = calculator
 
     if isinstance(filter, str):
         if filter not in _valid_filters:

mlip_arena/tasks/phonon.py

@@ -97,11 +97,9 @@ def _generate_task_run_name():
     parameters = task_run.parameters
 
     atoms = parameters["atoms"]
-    calculator = parameters["calculator"]
+    calculator_name = parameters["calculator"]
 
-    return (
-        f"{task_name}: {atoms.get_chemical_formula()} - {calculator.__class__.__name__}"
-    )
+    return f"{task_name}: {atoms.get_chemical_formula()} - {calculator_name}"
 
 
 @task(
@@ -124,7 +122,7 @@ def run(
     outdir: str | None = None,
 ):
     phonon = get_phonopy(
-        atoms=atoms,
+        atoms=atoms.copy(),
         supercell_matrix=supercell_matrix,
         min_lengths=min_lengths,
         symprec=symprec,

mlip_arena/tasks/utils.py

@@ -5,11 +5,11 @@ from __future__ import annotations
 from pprint import pformat
 
 import torch
-from torch_dftd.torch_dftd3_calculator import TorchDFTD3Calculator
-
 from ase import units
 from ase.calculators.calculator import BaseCalculator
 from ase.calculators.mixing import SumCalculator
+from torch_dftd.torch_dftd3_calculator import TorchDFTD3Calculator
+
 from mlip_arena.models import MLIPEnum
 
 try:
@@ -72,6 +72,7 @@ def get_calculator(
 
     if isinstance(calculator_name, MLIPEnum) and calculator_name in MLIPEnum:
         calc = calculator_name.value(**calculator_kwargs)
+        calc.__str__ = lambda: calculator_name.name
     elif isinstance(calculator_name, str) and hasattr(MLIPEnum, calculator_name):
         calc = MLIPEnum[calculator_name].value(**calculator_kwargs)
     elif isinstance(calculator_name, type) and issubclass(
@@ -79,11 +80,13 @@ def get_calculator(
     ):
         logger.warning(f"Using custom calculator class: {calculator_name}")
         calc = calculator_name(**calculator_kwargs)
+        calc.__str__ = lambda: f"{calc.__class__.__name__}"
     elif isinstance(calculator_name, BaseCalculator):
         logger.warning(
             f"Using custom calculator object (kwargs are ignored): {calculator_name}"
         )
         calc = calculator_name
+        calc.__str__ = lambda: f"{calc.__class__.__name__}"
     else:
         raise ValueError(f"Invalid calculator: {calculator_name}")
 
@@ -107,5 +110,5 @@ def get_calculator(
         if dispersion_kwargs:
             logger.info(pformat(dispersion_kwargs))
 
-    assert isinstance(calc, BaseCalculator)
+    assert isinstance(calc, BaseCalculator)    
     return calc

tests/test_elasticity.py

@@ -4,6 +4,7 @@ import numpy as np
 import pytest
 from mlip_arena.models import MLIPEnum
 from mlip_arena.tasks.elasticity import run as ELASTICITY
+from mlip_arena.tasks.utils import get_calculator
 from prefect.testing.utilities import prefect_test_harness
 
 from ase.build import bulk
@@ -22,9 +23,9 @@ def test_elasticity(model: MLIPEnum):
     with prefect_test_harness():
         result = ELASTICITY(
             atoms=bulk("Cu", "fcc", a=3.6),
-            calculator_name=model.name,
-            calculator_kwargs={},
-            device=None,
+            calculator=get_calculator(
+                calculator_name=model.name,
+            ),
             optimizer="BFGSLineSearch",
             optimizer_kwargs=None,
             filter="FrechetCell",

tests/test_eos.py

@@ -7,6 +7,8 @@ from prefect.testing.utilities import prefect_test_harness
 
 from mlip_arena.models import MLIPEnum
 from mlip_arena.tasks.eos import run as EOS
+from mlip_arena.tasks.utils import get_calculator
+
 
 
 @flow(persist_result=True)
@@ -17,9 +19,9 @@ def single_eos_flow(calculator_name, concurrent=True, cache=False):
         refresh_cache=not cache,
     )(
         atoms=atoms,
-        calculator_name=calculator_name,
-        calculator_kwargs={},
-        device=None,
+        calculator=get_calculator(
+            calculator_name=calculator_name,
+        ),
         optimizer="BFGSLineSearch",
         optimizer_kwargs=None,
         filter="FrechetCell",
@@ -62,4 +64,4 @@ def test_eos(model: MLIPEnum, concurrent: bool):
             cache=True,
         )
         assert isinstance(b0_cache := result["b0"], float)
-        assert b0_scratch == pytest.approx(b0_cache, rel=1e-6)
+        assert b0_scratch == pytest.approx(b0_cache, rel=1e-5)

tests/test_md.py

@@ -6,6 +6,7 @@ from ase.build import bulk
 
 from mlip_arena.models import MLIPEnum
 from mlip_arena.tasks.md import run as MD
+from mlip_arena.tasks.utils import get_calculator
 
 atoms = bulk("Cu", "fcc", a=3.6)
 
@@ -15,8 +16,9 @@ def test_nve(model: MLIPEnum):
 
     result = MD.fn(
         atoms,
-        calculator_name=model.name,
-        calculator_kwargs={},
+        calculator=get_calculator(
+            calculator_name=model.name,
+        ),
         ensemble="nve",
         dynamics="velocityverlet",
         total_time=10,

tests/test_neb.py

@@ -3,6 +3,7 @@ import sys
 import pytest
 from mlip_arena.models import MLIPEnum
 from mlip_arena.tasks import NEB_FROM_ENDPOINTS
+from mlip_arena.tasks.utils import get_calculator
 from prefect.testing.utilities import prefect_test_harness
 
 from ase.spacegroup import crystal
@@ -35,7 +36,9 @@ def test_neb(model: MLIPEnum):
             start=start.copy(),
             end=end.copy(),
             n_images=5,
-            calculator_name=model.name,
+            calculator=get_calculator(
+                calculator_name=model.name,
+            ),
             optimizer="FIRE2",
         )