Skip to main content
Tasks are the core building blocks of MLIP Arena. Each task wraps a well-defined atomistic simulation — such as structure optimization or molecular dynamics — as a Prefect task, giving you automatic caching, parallel execution, and composability with no extra boilerplate.

Available tasks

Structure optimization

Relax atomic positions and/or cell parameters to a local energy minimum.

Equation of state

Compute the energy-volume curve and fit a Birch-Murnaghan EOS to extract bulk modulus.

Molecular dynamics

Run NVE, NVT, or NPT simulations with flexible temperature and pressure schedules.

Phonons

Calculate phonon band structures, DOS, and thermal properties via phonopy.

Nudged elastic band

Find minimum-energy paths and transition-state barriers between two structures.

Elasticity

Compute the full elastic tensor and derived moduli from finite-strain deformations.

Importing tasks

All tasks are importable from the top-level mlip_arena.tasks package:
from mlip_arena.tasks import OPT, EOS, MD, PHONON, NEB, NEB_FROM_ENDPOINTS, ELASTICITY
Alternatively, import from each sub-module:
from mlip_arena.tasks.optimize import run as OPT
from mlip_arena.tasks.eos     import run as EOS
from mlip_arena.tasks.md      import run as MD
from mlip_arena.tasks.phonon  import run as PHONON
from mlip_arena.tasks.neb     import run as NEB, run_from_endpoints as NEB_FROM_ENDPOINTS
from mlip_arena.tasks.elasticity import run as ELASTICITY
PHONON requires phonopy. If phonopy is not installed, the import falls back gracefully and logs a warning.

Caching and Prefect integration

Every task is decorated with @task(cache_policy=TASK_SOURCE + INPUTS). This means:
  • Automatic caching — if you call a task with the same inputs and the task source code has not changed, Prefect returns the cached result instantly.
  • Parallel execution — call .submit() instead of calling the task directly to dispatch it as a non-blocking future.
  • Composability — tasks can call other tasks. For example, EOS internally submits multiple OPT tasks concurrently.

General usage pattern

Every task follows the same pattern:
1

Build or load an ASE Atoms object

from ase.build import bulk
atoms = bulk("Cu", "fcc", a=3.6)
2

Instantiate a calculator

Use get_calculator from mlip_arena.tasks.utils to load any registered MLIP, or pass any ASE BaseCalculator directly.
from mlip_arena.models import MLIPEnum
from mlip_arena.tasks.utils import get_calculator

calc = get_calculator(MLIPEnum.MACE_MP)
3

Call the task

result = OPT(atoms=atoms, calculator=calc)
print(result["atoms"].get_potential_energy())
4

(Optional) Run in parallel inside a flow

from prefect import flow
from mlip_arena.tasks import MD

@flow
def benchmark():
    futures = [
        MD.submit(atoms=atoms, calculator=get_calculator(m))
        for m in MLIPEnum
    ]
    return [f.result(raise_on_failure=False) for f in futures]

benchmark()

Dispersion corrections

All tasks accept a pre-built ASE calculator, so dispersion corrections (e.g. DFT-D3 via torch_dftd) can be applied by composing calculators before passing them in:
from mlip_arena.tasks.utils import get_calculator
from ase import units

calc = get_calculator(
    "MACE_MP",
    dispersion=True,
    dispersion_kwargs=dict(damping="bj", xc="pbe", cutoff=40.0 * units.Bohr),
)