From 5862e3513b03a8e86c4e3afbb8ac201bf99d5697 Mon Sep 17 00:00:00 2001 From: Marvin Poul Date: Fri, 3 Apr 2026 17:41:45 -0400 Subject: [PATCH 1/3] Make PyACECalculator deal with pickle The default reducer puts out the evaluator and a bunch of other C++ extension types that pickle then chokes on. Side step the problem here by providing a reducer that just makes unpickling equivalent to a new instantiation. The only difficulty here is that basis_set stays a file name if initialized like that and that PyACECalculator.__init__ expects it as a positional argument. Avoid the filename by dumping the whole bbasis config, so that unpickling can happen on other machines or in different locations. --- src/pyace/asecalc.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/pyace/asecalc.py b/src/pyace/asecalc.py index 5228759..30c760f 100644 --- a/src/pyace/asecalc.py +++ b/src/pyace/asecalc.py @@ -212,6 +212,10 @@ def dump_current_configuration(self, atoms, max_gamma): write(fname, atoms, format="cfg") self.current_extrapolation_structure_index += 1 + def __reduce__(self): + di = self.todict() + di.pop('basis_set', None) + return (type(self), (self.basis.to_BBasisConfiguration(),), self.todict()) class PyACEEnsembleCalculator(Calculator): """ From cafa43fb67e274351ada4369b9373008a7188c8f Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 15 Jul 2026 16:49:37 +0000 Subject: [PATCH 2/3] Support pickling of calculators with Ctilde basis sets ACECTildeBasisSet cannot be converted back to a BBasisConfiguration, so the previous reducer failed for calculators loaded from .ace/.yace potentials. Pass the Ctilde basis set through as is instead; its pybind11 binding already provides native pickle support. While at it, preserve the constructor kwargs (recursive_evaluator, gamma bounds, etc.) across the pickle round trip via a module-level helper, and drop the stale todict() state that previously leaked the original basis_set parameter (e.g. a file name) into the unpickled instance. Add round-trip tests for B-basis, Ctilde and recursive Ctilde calculators. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01BHj67uGNHN5bMH7LPoCR1J --- src/pyace/asecalc.py | 18 +++++++++++--- tests/test_PyACECalculator.py | 44 +++++++++++++++++++++++++++++++++++ 2 files changed, 59 insertions(+), 3 deletions(-) diff --git a/src/pyace/asecalc.py b/src/pyace/asecalc.py index 30c760f..d9cb20b 100644 --- a/src/pyace/asecalc.py +++ b/src/pyace/asecalc.py @@ -213,9 +213,21 @@ def dump_current_configuration(self, atoms, max_gamma): self.current_extrapolation_structure_index += 1 def __reduce__(self): - di = self.todict() - di.pop('basis_set', None) - return (type(self), (self.basis.to_BBasisConfiguration(),), self.todict()) + # 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 d0100e9..ea2dcc4 100644 --- a/tests/test_PyACECalculator.py +++ b/tests/test_PyACECalculator.py @@ -1,3 +1,5 @@ +import pickle + import numpy as np import pytest @@ -258,3 +260,45 @@ def check(at, msg): at = Atoms("H2", positions=[[0, 0, 0], [0, 0, 1.5]]) # ZBL check(at, "ZBL forces are inconsistent") + + +def check_pickle_roundtrip(calc): + a = create_dimer(2.0) + a.set_calculator(calc) + e1 = a.get_potential_energy() + f1 = a.get_forces() + + calc2 = pickle.loads(pickle.dumps(calc)) + + b = create_dimer(2.0) + b.set_calculator(calc2) + e2 = b.get_potential_energy() + f2 = b.get_forces() + + 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_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) From a338bfc1f55f78f5e4635e20e94b18cea3058806 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 15 Jul 2026 17:13:23 +0000 Subject: [PATCH 3/3] Test pickling with .yace potential and dimer minimization Address review: cover the Al-Ni .yace test potential and compare energies/forces along a whole BFGS dimer minimization trajectory between the original and the unpickled calculator. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01BHj67uGNHN5bMH7LPoCR1J --- tests/test_PyACECalculator.py | 39 +++++++++++++++++++++++++++-------- 1 file changed, 30 insertions(+), 9 deletions(-) diff --git a/tests/test_PyACECalculator.py b/tests/test_PyACECalculator.py index ea2dcc4..4aee55d 100644 --- a/tests/test_PyACECalculator.py +++ b/tests/test_PyACECalculator.py @@ -262,19 +262,33 @@ def check(at, msg): check(at, "ZBL forces are inconsistent") -def check_pickle_roundtrip(calc): - a = create_dimer(2.0) - a.set_calculator(calc) - e1 = a.get_potential_energy() - f1 = a.get_forces() +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)) - b = create_dimer(2.0) - b.set_calculator(calc2) - e2 = b.get_potential_energy() - f2 = b.get_forces() + # 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 @@ -294,6 +308,13 @@ def test_pickle_ctilde_calculator(): 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