Tasks overview

Modular, Prefect-powered computational tasks for benchmarking MLIPs.

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),
)