cache ase calculator

Author: esoteric-ephemera

src/atomate2/ase/jobs.py

@@ -54,8 +54,7 @@ class AseMaker(Maker, ABC):
     class EMTStaticMaker(AseMaker):
         name: str = "EMT static maker"
 
-        @property
-        def calculator(self):
+        def _get_calculator(self):
             return EMT()
     ```
 
@@ -95,6 +94,10 @@ def calculator(self):
     store_trajectory: StoreTrajectoryOption = StoreTrajectoryOption.NO
     tags: list[str] | None = None
 
+    def __post_init__(self) -> None:
+        """Enable caching of the ASE calculator via private attribute."""
+        self._calculator: Calculator | None = None
+
     @job(data=_ASE_DATA_OBJECTS)
     def make(
         self,
@@ -148,11 +151,18 @@ def run_ase(
             elapsed_time=t_f - t_i,
         )
 
-    @property
     @abstractmethod
+    def _get_calculator(self) -> Calculator:
+        """Load ASE calculator, to be implemented by the user."""
+
+    @property
     def calculator(self) -> Calculator:
-        """ASE calculator, method to be implemented in subclasses."""
-        raise NotImplementedError
+        """Retrieve cached ASE calculator."""
+        if getattr(self, "_calculator", None) is None:
+            self._calculator = self._get_calculator()
+        if self._calculator is None:
+            raise ValueError("ASE calculator not properly initialized.")
+        return self._calculator
 
 
 @dataclass
@@ -208,8 +218,7 @@ class AseRelaxMaker(AseMaker):
 
     def __post_init__(self) -> None:
         """Ensure that physical relaxation settings are used."""
-        if hasattr(super(), "__post_init__"):
-            super().__post_init__()  # type: ignore[misc]
+        super().__post_init__()
         if self.relax_cell and self.relax_shape:
             raise ValueError(
                 "You have set both `relax_cell` (relaxing the cell shape and volume) "
@@ -299,8 +308,7 @@ class EmtRelaxMaker(AseRelaxMaker):
 
     name: str = "EMT relaxation"
 
-    @property
-    def calculator(self) -> Calculator:
+    def _get_calculator(self) -> Calculator:
         """EMT calculator."""
         from ase.calculators.emt import EMT
 
@@ -320,8 +328,7 @@ class LennardJonesRelaxMaker(AseRelaxMaker):
 
     name: str = "Lennard-Jones 6-12 relaxation"
 
-    @property
-    def calculator(self) -> Calculator:
+    def _get_calculator(self) -> None:
         """Lennard-Jones calculator."""
         from ase.calculators.lj import LennardJones
 
@@ -378,8 +385,7 @@ class GFNxTBRelaxMaker(AseRelaxMaker):
         }
     )
 
-    @property
-    def calculator(self) -> Calculator:
+    def _get_calculator(self) -> None:
         """GFN-xTB / TBLite calculator."""
         try:
             from tblite.ase import TBLite

src/atomate2/ase/md.py

@@ -8,7 +8,7 @@
 import os
 import sys
 import time
-from abc import ABC, abstractmethod
+from abc import ABC
 from collections.abc import Sequence
 from dataclasses import dataclass, field
 from enum import Enum
@@ -189,6 +189,7 @@ class AseMDMaker(AseMaker, ABC):
 
     def __post_init__(self) -> None:
         """Ensure that ensemble is an enum."""
+        super().__post_init__()
         if isinstance(self.ensemble, str):
             self.ensemble = MDEnsemble(self.ensemble.split("MDEnsemble.")[-1])
 
@@ -444,12 +445,6 @@ def _callback(dyn: MolecularDynamics = md_runner) -> None:
             elapsed_time=t_f - t_i,
         )
 
-    @property
-    @abstractmethod
-    def calculator(self) -> Calculator:
-        """ASE calculator, to be overwritten by user."""
-        raise NotImplementedError
-
 
 @dataclass
 class LennardJonesMDMaker(AseMDMaker):
@@ -461,8 +456,7 @@ class LennardJonesMDMaker(AseMDMaker):
 
     name: str = "Lennard-Jones 6-12 MD"
 
-    @property
-    def calculator(self) -> Calculator:
+    def _get_calculator(self) -> Calculator:
         """Lennard-Jones calculator."""
         from ase.calculators.lj import LennardJones
 
@@ -495,8 +489,7 @@ class GFNxTBMDMaker(AseMDMaker):
         }
     )
 
-    @property
-    def calculator(self) -> Calculator:
+    def _get_calculator(self) -> Calculator:
         """GFN-xTB / TBLite calculator."""
         try:
             from tblite.ase import TBLite

src/atomate2/ase/neb.py

@@ -257,8 +257,7 @@ class EmtNebFromImagesMaker(AseNebFromImagesMaker):
 
     name: str = "EMT NEB from images maker"
 
-    @property
-    def calculator(self) -> Calculator:
+    def _get_calculator(self) -> Calculator:
         """EMT calculator."""
         from ase.calculators.emt import EMT
 

src/atomate2/forcefields/utils.py

@@ -210,8 +210,7 @@ def _run_ase_safe(self, *args, **kwargs) -> AseResult:
         with revert_default_dtype():
             return self.run_ase(*args, **kwargs)
 
-    @property
-    def calculator(self) -> Calculator:
+    def _get_calculator(self) -> Calculator:
         """ASE calculator, can be overwritten by user."""
         return ase_calculator(
             self.calculator_meta,

tests/ase/test_jobs.py

@@ -30,17 +30,15 @@
 class EMTStaticMaker(AseMaker):
     name: str = "EMT static maker"
 
-    @property
-    def calculator(self):
+    def _get_calculator(self):
         return EMT()
 
 
 @dataclass
 class EMTRelaxMaker(AseRelaxMaker):
     name: str = "EMT relax maker"
 
-    @property
-    def calculator(self):
+    def _get_calculator(self):
         return EMT()
 
 

tests/ase/test_neb.py

@@ -47,8 +47,7 @@ class EmtNebFromEndpointsMaker(AseNebFromEndpointsMaker):
         default_factory=EmtRelaxMaker,
     )
 
-    @property
-    def calculator(self):
+    def _get_calculator(self):
         return EMT(**self.calculator_kwargs)