promote tasks up to module level

mlip_arena/tasks/__init__.py

@@ -6,7 +6,17 @@ from huggingface_hub import HfApi, HfFileSystem, hf_hub_download
 from mlip_arena.models import MLIP
 from mlip_arena.models import REGISTRY as MODEL_REGISTRY
 
-with open(Path(__file__).parent / "registry.yaml") as f:
+
+from .elasticity import run as ELASTICITY
+from .eos import run as EOS
+from .md import run as MD
+from .neb import run as NEB
+from .neb import run_from_endpoints as NEB_FROM_ENDPOINTS
+from .optimize import run as OPT
+
+__all__ = ["OPT", "EOS", "MD", "NEB", "NEB_FROM_ENDPOINTS", "ELASTICITY"]
+
+with open(Path(__file__).parent / "registry.yaml", encoding="utf-8") as f:
     REGISTRY = yaml.safe_load(f)
 
 

mlip_arena/tasks/md.py

@@ -202,7 +202,7 @@ def _generate_task_run_name():
 def run(
     atoms: Atoms,
     calculator_name: str | MLIPEnum,
-    calculator_kwargs: dict | None,
+    calculator_kwargs: dict | None = None,
     dispersion: bool = False,
     dispersion_kwargs: dict | None = None,
     device: str | None = None,

mlip_arena/tasks/neb.py

@@ -167,11 +167,11 @@ def run(
 
 
 @task(
-    name="NEB from end points",
+    name="NEB from endpoints",
     task_run_name=_generate_task_run_name,
     cache_policy=TASK_SOURCE + INPUTS,
 )
-def run_from_end_points(
+def run_from_endpoints(
     start: Atoms,
     end: Atoms,
     n_images: int,

tests/test_neb.py

@@ -2,8 +2,7 @@ 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 mlip_arena.tasks import NEB_FROM_ENDPOINTS
 from prefect.testing.utilities import prefect_test_harness
 
 from ase.spacegroup import crystal
@@ -32,7 +31,7 @@ def test_neb(model: MLIPEnum):
     """
 
     with prefect_test_harness():
-        result = NEB(
+        result = NEB_FROM_ENDPOINTS(
             start=start.copy(),
             end=end.copy(),
             n_images=5,