improve caching, unify calculator init

mlip_arena/models/externals/mace-mp.py

@@ -11,7 +11,7 @@ from mlip_arena.models.utils import get_freer_device
 class MACE_MP_Medium(MACECalculator):
     def __init__(
         self,
-        checkpoint="http://tinyurl.com/5yyxdm76",
+        checkpoint="https://github.com/ACEsuit/mace-mp/releases/download/mace_mp_0/2023-12-03-mace-128-L1_epoch-199.model",
         device: str | None = None,
         default_dtype="float32",
         **kwargs,

mlip_arena/models/externals/mattersim.py

@@ -5,11 +5,14 @@ from pathlib import Path
 import yaml
 from mattersim.forcefield import MatterSimCalculator
 
+from ase import Atoms
 from mlip_arena.models.utils import get_freer_device
+# from pymatgen.io.ase import AseAtomsAdaptor, MSONAtoms
 
 with open(Path(__file__).parents[1] / "registry.yaml", encoding="utf-8") as f:
     REGISTRY = yaml.safe_load(f)
 
+
 class MatterSim(MatterSimCalculator):
     def __init__(
         self,
@@ -18,7 +21,19 @@ class MatterSim(MatterSimCalculator):
         **kwargs,
     ):
         super().__init__(
-            load_path=checkpoint, 
-            device=str(device or get_freer_device()), 
-            **kwargs
-        )
+            load_path=checkpoint, device=str(device or get_freer_device()), **kwargs
+        )
+
+    def calculate(
+        self,
+        atoms: Atoms | None = None,
+        properties: list | None = None,
+        system_changes: list | None = None,
+    ):
+        super().calculate(atoms, properties, system_changes)
+
+        # # convert unpicklizable atoms back to picklizable atoms to avoid prefect pickling error
+        # if isinstance(self.atoms, MSONAtoms):
+        #     atoms = self.atoms.copy()
+        #     strucutre = AseAtomsAdaptor().get_structure(atoms)
+        #     self.atoms = AseAtomsAdaptor().get_atoms(strucutre, msonable=False)

mlip_arena/tasks/elasticity.py

@@ -90,7 +90,7 @@ def run(
     normal_strains: list[float] | np.ndarray | None = np.linspace(-0.01, 0.01, 4),
     shear_strains: list[float] | np.ndarray | None = np.linspace(-0.06, 0.06, 4),
     persist_opt: bool = True,
-    cache_opt: bool = True,
+    cache_opt: bool = False,
 ) -> dict[str, Any] | State:
     """
     Compute the elastic tensor for the given structure and calculator.

mlip_arena/tasks/eos.py

@@ -17,8 +17,6 @@ from prefect.runtime import task_run
 from prefect.states import State
 
 from ase import Atoms
-from ase.filters import *  # type: ignore
-from ase.optimize import *  # type: ignore
 from ase.optimize.optimize import Optimizer
 from mlip_arena.models import MLIPEnum
 from mlip_arena.tasks.optimize import run as OPT
@@ -54,6 +52,7 @@ def run(
     max_abs_strain: float = 0.1,
     npoints: int = 11,
     concurrent: bool = True,
+    cache_opt: bool = False,
 ) -> dict[str, Any] | State:
     """
     Compute the equation of state (EOS) for the given atoms and calculator.
@@ -78,7 +77,12 @@ def run(
         A dictionary containing the EOS data, bulk modulus, equilibrium volume, and equilibrium energy if successful. Otherwise, a prefect state object.
     """
 
-    state = OPT(
+    OPT_ = OPT.with_options(
+        refresh_cache=not cache_opt,
+        persist_result=cache_opt,
+    )
+
+    state = OPT_(
         atoms=atoms,
         calculator_name=calculator_name,
         calculator_kwargs=calculator_kwargs,
@@ -112,7 +116,7 @@ def run(
             atoms = relaxed.copy()
             atoms.set_cell(c0 * f, scale_atoms=True)
 
-            future = OPT.submit(
+            future = OPT_.submit(
                 atoms=atoms,
                 calculator_name=calculator_name,
                 calculator_kwargs=calculator_kwargs,
@@ -138,7 +142,7 @@ def run(
             atoms = relaxed.copy()
             atoms.set_cell(c0 * f, scale_atoms=True)
 
-            state = OPT(
+            state = OPT_(
                 atoms=atoms,
                 calculator_name=calculator_name,
                 calculator_kwargs=calculator_kwargs,

mlip_arena/tasks/md.py

@@ -65,12 +65,9 @@ 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 torch_dftd.torch_dftd3_calculator import TorchDFTD3Calculator
 from tqdm.auto import tqdm
 
 from ase import Atoms, units
-from ase.calculators.calculator import Calculator
-from ase.calculators.mixing import SumCalculator
 from ase.io import read
 from ase.io.trajectory import Trajectory
 from ase.md.andersen import Andersen
@@ -86,7 +83,7 @@ from ase.md.velocitydistribution import (
 )
 from ase.md.verlet import VelocityVerlet
 from mlip_arena.models import MLIPEnum
-from mlip_arena.models.utils import get_freer_device
+from mlip_arena.tasks.utils import get_calculator
 
 _valid_dynamics: dict[str, tuple[str, ...]] = {
     "nve": ("velocityverlet",),
@@ -201,14 +198,12 @@ def _generate_task_run_name():
     name="MD",
     task_run_name=_generate_task_run_name,
     cache_policy=TASK_SOURCE + INPUTS
-    # cache_key_fn=task_input_hash,
-    # cache_expiration=timedelta(days=1)
 )
 def run(
     atoms: Atoms,
     calculator_name: str | MLIPEnum,
     calculator_kwargs: dict | None,
-    dispersion: str | None = None,
+    dispersion: bool = False,
     dispersion_kwargs: dict | None = None,
     device: str | None = None,
     ensemble: Literal["nve", "nvt", "npt"] = "nvt",
@@ -225,37 +220,14 @@ def run(
     traj_interval: int = 1,
     restart: bool = True,
 ):
-    device = device or str(get_freer_device())
-
-    print(f"Using device: {device}")
-
-    calculator_kwargs = calculator_kwargs or {}
-
-    if isinstance(calculator_name, MLIPEnum) and calculator_name in MLIPEnum:
-        assert issubclass(calculator_name.value, Calculator)
-        calc = calculator_name.value(**calculator_kwargs)
-    elif (
-        isinstance(calculator_name, str) and calculator_name in MLIPEnum._member_names_
-    ):
-        calc = MLIPEnum[calculator_name].value(**calculator_kwargs)
-    else:
-        raise ValueError(f"Invalid calculator: {calculator_name}")
-
-    print(f"Using calculator: {calc}")
-
-    dispersion_kwargs = dispersion_kwargs or {}
-
-    dispersion_kwargs.update({"device": device})
-
-    if dispersion is not None:
-        disp_calc = TorchDFTD3Calculator(
-            **dispersion_kwargs,
-        )
-        calc = SumCalculator([calc, disp_calc])
-
-        print(f"Using dispersion: {dispersion}")
-
-    atoms.calc = calc
+    
+    atoms.calc = get_calculator(
+        calculator_name=calculator_name,
+        calculator_kwargs=calculator_kwargs,
+        dispersion=dispersion,
+        dispersion_kwargs=dispersion_kwargs,
+        device=device,
+    )
 
     if time_step is None:
         # If a structure contains an isotope of hydrogen, set default `time_step`

mlip_arena/tasks/neb.py

@@ -57,6 +57,7 @@ from mlip_arena.tasks.optimize import run as OPT
 from mlip_arena.tasks.utils import get_calculator
 from pymatgen.io.ase import AseAtomsAdaptor
 
+
 if TYPE_CHECKING:
     pass
 
@@ -100,7 +101,7 @@ def run(
     images: list[Atoms],
     calculator_name: str | MLIPEnum,
     calculator_kwargs: dict | None = None,
-    dispersion: str | None = None,
+    dispersion: bool = False,
     dispersion_kwargs: dict | None = None,
     device: str | None = None,
     optimizer: Optimizer | str = "MDMin",  # type: ignore
@@ -159,11 +160,16 @@ def run(
     optimizer_instance.run(**criterion)
 
     neb_tool = NEBTools(neb.images)
+    barrier = neb_tool.get_barrier()
+
+    forcefit = fit_images(neb.images)
+
+    images = neb.images
 
     return {
-        "barrier": neb_tool.get_barrier(),
-        "images": neb.images,
-        "forcefit": fit_images(neb.images),
+        "barrier": barrier,
+        "images": images,
+        "forcefit": forcefit,
     }
 
 
@@ -188,6 +194,7 @@ def run_from_end_points(
     interpolation: Literal["linear", "idpp"] = "idpp",
     climb: bool = True,
     traj_file: str | Path | None = None,
+    cache_subtasks: bool = False,
 ) -> dict[str, Any] | State:
     """Run the nudged elastic band (NEB) calculation from end points.
 
@@ -212,7 +219,9 @@ def run_from_end_points(
     """
 
     if relax_end_points:
-        relax = OPT(
+        relax = OPT.with_options(
+            refresh_cache=not cache_subtasks,
+        )(
             atoms=start.copy(),
             calculator_name=calculator_name,
             calculator_kwargs=calculator_kwargs,
@@ -225,7 +234,9 @@ def run_from_end_points(
         )
         start = relax["atoms"]
 
-        relax = OPT(
+        relax = OPT.with_options(
+            refresh_cache=not cache_subtasks,
+        )(
             atoms=end.copy(),
             calculator_name=calculator_name,
             calculator_kwargs=calculator_kwargs,
@@ -252,7 +263,9 @@ def run_from_end_points(
 
     images = [s.to_ase_atoms() for s in path]
 
-    return run(
+    return run.with_options(
+        refresh_cache=not cache_subtasks,
+    )(
         images,
         calculator_name,
         calculator_kwargs=calculator_kwargs,

mlip_arena/tasks/optimize.py

@@ -7,18 +7,15 @@ from __future__ import annotations
 from prefect import task
 from prefect.cache_policies import INPUTS, TASK_SOURCE
 from prefect.runtime import task_run
-from torch_dftd.torch_dftd3_calculator import TorchDFTD3Calculator
 
 from ase import Atoms
-from ase.calculators.calculator import Calculator
-from ase.calculators.mixing import SumCalculator
 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.models.utils import get_freer_device
+from mlip_arena.tasks.utils import get_calculator
 
 _valid_filters: dict[str, Filter] = {
     "Filter": Filter,
@@ -54,17 +51,13 @@ def _generate_task_run_name():
 
 
 @task(
-    name="OPT",
-    task_run_name=_generate_task_run_name,
-    cache_policy=TASK_SOURCE + INPUTS
-    # cache_key_fn=task_input_hash,
-    # cache_expiration=timedelta(days=1)
+    name="OPT", task_run_name=_generate_task_run_name, cache_policy=TASK_SOURCE + INPUTS
 )
 def run(
     atoms: Atoms,
     calculator_name: str | MLIPEnum,
     calculator_kwargs: dict | None = None,
-    dispersion: str | None = None,
+    dispersion: bool = False,
     dispersion_kwargs: dict | None = None,
     device: str | None = None,
     optimizer: Optimizer | str = BFGSLineSearch,
@@ -74,37 +67,13 @@ def run(
     criterion: dict | None = None,
     symmetry: bool = False,
 ):
-    device = device or str(get_freer_device())
-
-    print(f"Using device: {device}")
-
-    calculator_kwargs = calculator_kwargs or {}
-
-    if isinstance(calculator_name, MLIPEnum) and calculator_name in MLIPEnum:
-        assert issubclass(calculator_name.value, Calculator)
-        calc = calculator_name.value(**calculator_kwargs)
-    elif (
-        isinstance(calculator_name, str) and calculator_name in MLIPEnum._member_names_
-    ):
-        calc = MLIPEnum[calculator_name].value(**calculator_kwargs)
-    else:
-        raise ValueError(f"Invalid calculator: {calculator_name}")
-
-    print(f"Using calculator: {calc}")
-
-    dispersion_kwargs = dispersion_kwargs or {}
-
-    dispersion_kwargs.update({"device": device})
-
-    if dispersion is not None:
-        disp_calc = TorchDFTD3Calculator(
-            **dispersion_kwargs,
-        )
-        calc = SumCalculator([calc, disp_calc])
-
-        print(f"Using dispersion: {dispersion}")
-
-    atoms.calc = calc
+    atoms.calc = get_calculator(
+        calculator_name=calculator_name,
+        calculator_kwargs=calculator_kwargs,
+        dispersion=dispersion,
+        dispersion_kwargs=dispersion_kwargs,
+        device=device,
+    )
 
     if isinstance(filter, str):
         if filter not in _valid_filters:

mlip_arena/tasks/utils.py

@@ -4,24 +4,33 @@ from __future__ import annotations
 
 from torch_dftd.torch_dftd3_calculator import TorchDFTD3Calculator
 
+from ase import units
 from ase.calculators.calculator import Calculator
 from ase.calculators.mixing import SumCalculator
-from ase.filters import *  # type: ignore
-from ase.optimize import *  # type: ignore
 from mlip_arena.models import MLIPEnum
 from mlip_arena.models.utils import get_freer_device
 
+try:
+    from prefect.logging import get_run_logger
+
+    logger = get_run_logger()
+except (ImportError, RuntimeError):
+    from loguru import logger
+
+from pprint import pformat
+
 
 def get_calculator(
     calculator_name: str | MLIPEnum,
     calculator_kwargs: dict | None,
-    dispersion: str | None = None,
+    dispersion: bool = False,
     dispersion_kwargs: dict | None = None,
     device: str | None = None,
-) -> Calculator:
+) -> Calculator | SumCalculator:
+    """Get a calculator with optional dispersion correction."""
     device = device or str(get_freer_device())
 
-    print(f"Using device: {device}")
+    logger.info(f"Using device: {device}")
 
     calculator_kwargs = calculator_kwargs or {}
 
@@ -33,19 +42,25 @@ def get_calculator(
     else:
         raise ValueError(f"Invalid calculator: {calculator_name}")
 
-    print(f"Using calculator: {calc}")
+    logger.info(f"Using calculator: {calc}")
+    if calculator_kwargs:
+        logger.info(pformat(calculator_kwargs))
 
-    dispersion_kwargs = dispersion_kwargs or {}
+    dispersion_kwargs = dispersion_kwargs or dict(
+        damping="bj", xc="pbe", cutoff=40.0 * units.Bohr
+    )
 
     dispersion_kwargs.update({"device": device})
 
-    if dispersion is not None:
+    if dispersion:
         disp_calc = TorchDFTD3Calculator(
             **dispersion_kwargs,
         )
         calc = SumCalculator([calc, disp_calc])
 
-        print(f"Using dispersion: {dispersion}")
+        logger.info(f"Using dispersion: {disp_calc}")
+        if dispersion_kwargs:
+            logger.info(pformat(dispersion_kwargs))
 
-    assert isinstance(calc, Calculator)
+    assert isinstance(calc, Calculator) or isinstance(calc, SumCalculator)
     return calc