Merge pull request #1036 from pyiron/pickle_atoms

Pickle atoms

pyiron_atomistics/atomistics/structure/atoms.py

@@ -220,6 +220,14 @@ def __init__(
         # Velocities were not handled at all during file writing
         self._velocities = None
 
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+
+    def __getstate__(self):
+        # Only necessary to support pickling in python <3.11
+        # https://docs.python.org/release/3.11.2/library/pickle.html#object.__getstate__
+        return self.__dict__
+
     @property
     def velocities(self):
         return self._velocities

pyiron_atomistics/atomistics/structure/periodic_table.py

@@ -89,6 +89,14 @@ def __getitem__(self, item):
         if item in self.sub.index:
             return self.sub[item]
 
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+
+    def __getstate__(self):
+        # Only necessary to support pickling in python <3.11
+        # https://docs.python.org/release/3.11.2/library/pickle.html#object.__getstate__
+        return self.__dict__
+
     def __eq__(self, other):
         if self is other:
             return True
@@ -223,6 +231,17 @@ def __getitem__(self, item):
         if item in self.dataframe.index.values:
             return self.dataframe.loc[item]
 
+    def __setstate__(self, state):
+        """
+        Used by (cloud)pickle; force the state update to avoid recursion pickling Atoms
+        """
+        self.__dict__.update(state)
+
+    def __getstate__(self):
+        # Only necessary to support pickling in python <3.11
+        # https://docs.python.org/release/3.11.2/library/pickle.html#object.__getstate__
+        return self.__dict__
+
     def from_hdf(self, hdf):
         """
         loads an element with his parameters from the hdf5 job file by creating an Object of the ChemicalElement type.

tests/atomistics/structure/test_atoms.py

@@ -5,6 +5,7 @@
 import unittest
 import numpy as np
 import os
+import pickle
 import time
 import warnings
 from pyiron_atomistics.atomistics.structure.atom import Atom
@@ -1953,6 +1954,12 @@ def test_calc_to_hdf(self):
                 msg=f"Calculator parameter {k} not correctly restored from HDF!",
             )
 
+    def test_pickle(self):
+        pickled = pickle.dumps(self.C3)
+        unpickled = pickle.loads(pickled)
+        self.assertEqual(unpickled, self.C3)
+        self.assertTrue(np.allclose(unpickled.cell, self.C3.cell))
+
 
 def generate_fcc_lattice(a=4.2):
     positions = [[0, 0, 0]]