Skip to content

Commit

Permalink
promote tasks up to module level
Browse files Browse the repository at this point in the history
  • Loading branch information
chiang-yuan committed Jan 12, 2025
1 parent e80e29d commit f4bf91c
Show file tree
Hide file tree
Showing 3 changed files with 15 additions and 6 deletions.
12 changes: 11 additions & 1 deletion mlip_arena/tasks/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,17 @@
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)


Expand Down
4 changes: 2 additions & 2 deletions mlip_arena/tasks/neb.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
5 changes: 2 additions & 3 deletions tests/test_neb.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,7 @@

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
Expand Down Expand Up @@ -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,
Expand Down

0 comments on commit f4bf91c

Please sign in to comment.