Allow caching of restart data in ASE calculator (#18)

Author: awvwgk

xtb/ase/calculator.py

@@ -51,6 +51,7 @@
  electronic_temperature   300.0        Electronic temperatur for TB methods
  max_iterations           250          Iterations for self-consistent evaluation
  solvent                  "none"       GBSA implicit solvent model
+ cache_api                True         Reuse generate API objects (recommended)
 ======================== ============ ============================================
 """
 
@@ -84,6 +85,7 @@ class XTB(ase_calc.Calculator):
         "max_iterations": 250,
         "electronic_temperature": 300.0,
         "solvent": "None",
+        "cache_api": True,
     }
 
     _res = None
@@ -96,49 +98,85 @@ def __init__(
 
         ase_calc.Calculator.__init__(self, atoms=atoms, **kwargs)
 
-        # loads the default parameters and updates with actual values
-        self.parameters = self.get_default_parameters()
-        # now set all parameters
-        self.set(**kwargs)
-
-    def set(self, **kwargs):
+    def set(self, **kwargs) -> dict:
         """Set new parameters to xtb"""
 
         changed_parameters = ase_calc.Calculator.set(self, **kwargs)
 
-        self._check(changed_parameters)
+        self._check_parameters(changed_parameters)
 
-        # Always reset the xtb calculator for now
+        # Always reset the calculation if parameters change
         if changed_parameters:
             self.reset()
 
+        # If the method is changed, invalidate the cached calculator as well
+        if "method" in changed_parameters:
+            self._xtb = None
+            self._res = None
+
+        # Minor changes can be updated in the API calculator directly
+        if self._xtb is not None:
+            if "accuracy" in changed_parameters:
+                self._xtb.set_accuracy(self.parameters.accuracy)
+
+            if "electronic_temperature" in changed_parameters:
+                self._xtb.set_electronic_temperature(
+                    self.parameters.electronic_temperature
+                )
+
+            if "max_iterations" in changed_parameters:
+                self._xtb.set_max_iterations(self.parameters.max_iterations)
+
+            if "solvent" in changed_parameters:
+                self._xtb.set_solvent(get_solvent(self.parameters.solvent))
+
         return changed_parameters
 
-    def _check(self, parameters):
+    def _check_parameters(self, parameters: dict) -> None:
         """Verifiy provided parameters are valid"""
 
         if "method" in parameters and get_method(parameters["method"]) is None:
             raise ase_calc.InputError(
                 "Invalid method {} provided".format(parameters["method"])
             )
 
-    def reset(self):
+    def reset(self) -> None:
         """Clear all information from old calculation"""
         ase_calc.Calculator.reset(self)
 
-        self._res = None
-
-    def calculate(
-        self,
-        atoms: Optional[Atoms] = None,
-        properties: List[str] = None,
-        system_changes: List[str] = ase_calc.all_changes,
-    ):
-        """Perform actual calculation with by calling the xtb API"""
-
-        if not properties:
-            properties = ["energy"]
-        ase_calc.Calculator.calculate(self, atoms, properties, system_changes)
+        if not self.parameters.cache_api:
+            self._xtb = None
+            self._res = None
+
+    def _check_api_calculator(self, system_changes: List[str]) -> None:
+        """Check state of API calculator and reset if necessary"""
+
+        # Changes in positions and cell parameters can use a normal update
+        _reset = system_changes.copy()
+        if "positions" in _reset:
+            _reset.remove("positions")
+        if "cell" in _reset:
+            _reset.remove("cell")
+
+        # Invalidate cached calculator and results object
+        if _reset:
+            self._xtb = None
+            self._res = None
+        else:
+            if system_changes and self._xtb is not None:
+                try:
+                    _cell = self.atoms.cell
+                    self._xtb.update(
+                        self.atoms.positions / Bohr, _cell / Bohr,
+                    )
+                # An exception in this part means the geometry is bad,
+                # still we will give a complete reset a try as well
+                except XTBException:
+                    self._xtb = None
+                    self._res = None
+
+    def _create_api_calculator(self) -> Calculator:
+        """Create a new API calculator object"""
 
         _method = get_method(self.parameters.method)
         if _method is None:
@@ -152,7 +190,7 @@ def calculate(
             _charge = self.atoms.get_initial_charges().sum()
             _uhf = int(self.atoms.get_initial_magnetic_moments().sum().round())
 
-            self._xtb = Calculator(
+            calc = Calculator(
                 _method,
                 self.atoms.numbers,
                 self.atoms.positions / Bohr,
@@ -161,15 +199,34 @@ def calculate(
                 _cell / Bohr,
                 _periodic,
             )
-            self._xtb.set_verbosity(VERBOSITY_MUTED)
-            self._xtb.set_accuracy(self.parameters.accuracy)
-            self._xtb.set_electronic_temperature(self.parameters.electronic_temperature)
-            self._xtb.set_max_iterations(self.parameters.max_iterations)
-            self._xtb.set_solvent(get_solvent(self.parameters.solvent))
+            calc.set_verbosity(VERBOSITY_MUTED)
+            calc.set_accuracy(self.parameters.accuracy)
+            calc.set_electronic_temperature(self.parameters.electronic_temperature)
+            calc.set_max_iterations(self.parameters.max_iterations)
+            calc.set_solvent(get_solvent(self.parameters.solvent))
 
         except XTBException:
             raise ase_calc.InputError("Cannot construct calculator for xtb")
 
+        return calc
+
+    def calculate(
+        self,
+        atoms: Optional[Atoms] = None,
+        properties: List[str] = None,
+        system_changes: List[str] = ase_calc.all_changes,
+    ) -> None:
+        """Perform actual calculation with by calling the xtb API"""
+
+        if not properties:
+            properties = ["energy"]
+        ase_calc.Calculator.calculate(self, atoms, properties, system_changes)
+
+        self._check_api_calculator(system_changes)
+
+        if self._xtb is None:
+            self._xtb = self._create_api_calculator()
+
         try:
             self._res = self._xtb.singlepoint(self._res)
         except XTBException:

xtb/ase/test_calculator.py

@@ -80,14 +80,22 @@ def test_gfn2_xtb_0d():
     ])
     dipole_moment = np.array([0.62120710, 0.28006659, 0.04465985])
 
-    calc = XTB(method="GFN2-xTB")
-    atoms.set_calculator(calc)
+    calc = XTB(method="GFN2-xTB", atoms=atoms)
 
     assert approx(atoms.get_potential_energy(), thr) == -592.6794366990786
     assert approx(atoms.get_forces(), thr) == forces
     assert approx(atoms.get_charges(), thr) == charges
     assert approx(atoms.get_dipole_moment(), thr) == dipole_moment
 
+    atoms.calc.set(
+        accuracy=0.1,
+        electronic_temperature=500.0,
+        max_iterations=20,
+        solvent="ch2cl2",
+    )
+
+    assert approx(atoms.get_potential_energy(), thr) == -592.9940608761889
+
 
 def test_gfn1_xtb_0d():
     """Test ASE interface to GFN1-xTB"""
@@ -140,8 +148,7 @@ def test_gfn1_xtb_0d():
     ])
     dipole_moment = np.array([0.76943477, 0.33021928, 0.05670150])
 
-    calc = XTB(method="GFN1-xTB")
-    atoms.set_calculator(calc)
+    atoms.calc = XTB(method="GFN1-xTB")
 
     assert approx(atoms.get_potential_energy(), thr) == -632.7363734598027
     assert approx(atoms.get_forces(), thr) == forces
@@ -235,6 +242,10 @@ def test_gfn2_xtb_3d():
     atoms.set_pbc(False)
     assert approx(atoms.get_potential_energy(), thr) == -1121.9196707084955
 
+    with raises(InputError):
+        atoms.positions = np.zeros((len(atoms), 3))
+        calc.calculate(atoms=atoms, system_changes=["positions"])
+
 
 def test_invalid_method():
     """GFN-xTB without method number is invalid, should raise an input error"""

xtb/ase/test_optimize.py

@@ -59,8 +59,7 @@ def test_gfn1xtb_bfgs():
         ]),
     )
 
-    calc = XTB(method="GFN1-xTB", accuracy=2.0)
-    atoms.set_calculator(calc)
+    atoms.calc = XTB(method="GFN1-xTB", accuracy=2.0, cache_api=False)
     opt = BFGS(atoms)
     opt.run(fmax=0.1)
 
@@ -107,7 +106,7 @@ def test_gfn2xtb_lbfgs():
     opt.run(fmax=0.1)
 
     assert approx(atoms.get_potential_energy(), thr) == -897.4533662470938
-    assert approx(np.linalg.norm(atoms.get_forces(), ord=2), thr) == 0.1939329480683042
+    assert approx(np.linalg.norm(atoms.get_forces(), ord=2), thr) == 0.19359647527783497
 
 
 def test_gfn2xtb_velocityverlet():
@@ -143,11 +142,17 @@ def test_gfn2xtb_velocityverlet():
         ]),
     )
 
-    calc = XTB(method="GFN2-xTB")
+    calc = XTB(method="GFN2-xTB", cache_api=False)
     atoms.set_calculator(calc)
 
     dyn = VelocityVerlet(atoms, timestep=1.0*fs)
     dyn.run(20)
 
     assert approx(atoms.get_potential_energy(), thr) == -896.9772346260584
     assert approx(atoms.get_kinetic_energy(), thr) == 0.022411127028842362
+
+    atoms.calc.set(cache_api=True)
+    dyn.run(20)
+
+    assert approx(atoms.get_potential_energy(), thr) == -896.9913862530841
+    assert approx(atoms.get_kinetic_energy(), thr) == 0.036580471363852810