diff --git a/src/pyace/asecalc.py b/src/pyace/asecalc.py index 359d6f9..2d19755 100644 --- a/src/pyace/asecalc.py +++ b/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): """ diff --git a/tests/test_PyACECalculator.py b/tests/test_PyACECalculator.py index 2f71236..74f7c1f 100644 --- a/tests/test_PyACECalculator.py +++ b/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)