Make PyACECalculator deal with pickle

src/pyace/asecalc.py

@@ -212,6 +212,22 @@ def dump_current_configuration(self, atoms, max_gamma):
         write(fname, atoms, format="cfg")
         self.current_extrapolation_structure_index += 1
 
+    def __reduce__(self):
+        # B-basis potentials are dumped as a full BBasisConfiguration, so that
+        # unpickling works on other machines or in different locations even if
+        # the calculator was created from a file name.  Ctilde basis sets
+        # cannot be converted back to a BBasisConfiguration, but their binding
+        # provides native pickle support, so they are passed through as is.
+        if isinstance(self.basis, ACEBBasisSet):
+            basis_spec = self.basis.to_BBasisConfiguration()
+        else:
+            basis_spec = self.basis
+        kwargs = {k: v for k, v in self.parameters.items() if k != "basis_set"}
+        return (_unpickle_calculator, (type(self), basis_spec, kwargs))
+
+
+def _unpickle_calculator(cls, basis_spec, kwargs):
+    return cls(basis_spec, **kwargs)
 
 class PyACEEnsembleCalculator(Calculator):
     """

tests/test_PyACECalculator.py

@@ -1,3 +1,5 @@
+import pickle
+
 import numpy as np
 import pytest
 
@@ -250,3 +252,66 @@ def check(at, msg):
 
     at = Atoms("H2", positions=[[0, 0, 0], [0, 0, 1.5]])  # ZBL
     check(at, "ZBL forces are inconsistent")
+
+
+def minimize_dimer(calc, elements):
+    from ase.optimize import BFGS
+
+    at = Atoms(elements, positions=[[0, 0, 0], [0, 0, 2.0]], pbc=False)
+    at.set_calculator(calc)
+
+    energies = [at.get_potential_energy()]
+    forces = [at.get_forces()]
+
+    opt = BFGS(at, logfile=None)
+    opt.attach(lambda: (energies.append(at.get_potential_energy()),
+                        forces.append(at.get_forces())))
+    opt.run(fmax=1e-3, steps=50)
+
+    return np.array(energies), np.array(forces)
+
+
+def check_pickle_roundtrip(calc, elements=("Al", "Al")):
+    calc2 = pickle.loads(pickle.dumps(calc))
+
+    # minimize the dimer with the original and the unpickled calculator and
+    # compare energies/forces along the whole trajectory
+    e1, f1 = minimize_dimer(calc, elements)
+    e2, f2 = minimize_dimer(calc2, elements)
+
+    assert len(e1) > 1
+    assert e1.shape == e2.shape
+    assert np.allclose(e1, e2)
+    assert np.allclose(f1, f2)
+    return calc2
+
+
+def test_pickle_bbasis_calculator():
+    calc = PyACECalculator(basis_set="tests/Al.pbe.13.2.yaml")
+    calc2 = check_pickle_roundtrip(calc)
+    assert isinstance(calc2.basis, ACEBBasisSet)
+    assert isinstance(calc2.evaluator, ACEBEvaluator)
+
+
+def test_pickle_ctilde_calculator():
+    calc = PyACECalculator(basis_set="tests/Al.pbe.rhocore.ace")
+    calc2 = check_pickle_roundtrip(calc)
+    assert isinstance(calc2.basis, ACECTildeBasisSet)
+    assert isinstance(calc2.evaluator, ACECTildeEvaluator)
+
+
+def test_pickle_yace_calculator():
+    calc = PyACECalculator(basis_set="tests/Al-Ni_opt_all.yace")
+    calc2 = check_pickle_roundtrip(calc, elements=("Al", "Ni"))
+    assert isinstance(calc2.basis, ACECTildeBasisSet)
+    assert isinstance(calc2.evaluator, ACECTildeEvaluator)
+
+
+def test_pickle_ctilde_recursive_calculator():
+    from pyace.evaluator import ACERecursiveEvaluator
+
+    calc = PyACECalculator(basis_set="tests/Al.pbe.rhocore.ace",
+                           recursive_evaluator=True, recursive=True)
+    calc2 = check_pickle_roundtrip(calc)
+    assert isinstance(calc2.basis, ACECTildeBasisSet)
+    assert isinstance(calc2.evaluator, ACERecursiveEvaluator)