From ca70d0d2ad91096bb720a22dd2a04291917dbbef 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 359d6f9..da63e66 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 0dda6cb0416903b1164a06b1bcd493d5ef05c45a 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 da63e66..2d19755 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 2f71236..efa27f3 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,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 650f717f9a0446a5ff0c8e994ce7b0b8bc615249 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 efa27f3..74f7c1f 100644 --- a/tests/test_PyACECalculator.py +++ b/tests/test_PyACECalculator.py @@ -254,19 +254,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 @@ -286,6 +300,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