diff --git a/src/python/espressomd/particle_data.pyx b/src/python/espressomd/particle_data.pyx index 5d3481cab00..791c4447b5e 100644 --- a/src/python/espressomd/particle_data.pyx +++ b/src/python/espressomd/particle_data.pyx @@ -56,6 +56,33 @@ cdef class ParticleHandle: cdef int update_particle_data(self) except -1: self.particle_data = &get_particle_data(self._id) + def to_dict(self): + """ + Returns the particle's attributes as a dictionary. + + It includes the content of ``particle_attributes``, minus a few exceptions: + + - :attr:`~ParticleHandle.dip`, :attr:`~ParticleHandle.director`: + Setting only the director will overwrite the orientation of the + particle around the axis parallel to dipole moment/director. + Quaternions contain the full info. + - :attr:`~ParticleHandle.image_box`, :attr:`~ParticleHandle.node` + + """ + + pickle_attr = copy(particle_attributes) + for i in ["director", "dip", "image_box", "node"]: + if i in pickle_attr: + pickle_attr.remove(i) + IF MASS == 0: + pickle_attr.remove("mass") + pdict = {} + + for property_ in pickle_attr: + pdict[property_] = ParticleHandle( + self.id).__getattribute__(property_) + return pdict + def __str__(self): res = collections.OrderedDict() # Id and pos first, then the rest @@ -1697,6 +1724,37 @@ class ParticleSlice(_ParticleSliceImpl): f"ParticleHandle does not have the attribute {name}.") super().__setattr__(name, value) + def to_dict(self): + """ + Returns the particles attributes as a dictionary. + + It can be used to save the particle data and recover it by using + + >>> p = system.part.add(...) + >>> particle_dict = p.to_dict() + >>> system.part.add(particle_dict) + + It includes the content of ``particle_attributes``, minus a few exceptions: + + - :attr:`~ParticleHandle.dip`, :attr:`~ParticleHandle.director`: + Setting only the director will overwrite the orientation of the + particle around the axis parallel to dipole moment/director. + Quaternions contain the full info. + - :attr:`~ParticleHandle.image_box`, :attr:`~ParticleHandle.node` + + """ + + odict = {} + key_list = [p.id for p in self] + for particle_number in key_list: + pdict = ParticleHandle(particle_number).to_dict() + for p_key, p_value in pdict.items(): + if p_key in odict: + odict[p_key].append(p_value) + else: + odict[p_key] = [p_value] + return odict + cdef class ParticleList: """ @@ -1733,21 +1791,11 @@ cdef class ParticleList: """ - pickle_attr = copy(particle_attributes) - for i in ["director", "dip", "id", "image_box", "node"]: - if i in pickle_attr: - pickle_attr.remove(i) - IF MASS == 0: - pickle_attr.remove("mass") odict = {} - key_list = [p.id for p in self] - for particle_number in key_list: - pdict = {} - - for property_ in pickle_attr: - pdict[property_] = ParticleHandle( - particle_number).__getattribute__(property_) - odict[particle_number] = pdict + for p in self: + pdict = p.to_dict() + del pdict["id"] + odict[p.id] = pdict return odict def __setstate__(self, params): diff --git a/testsuite/python/particle.py b/testsuite/python/particle.py index fb8236d740a..243e6d76d10 100644 --- a/testsuite/python/particle.py +++ b/testsuite/python/particle.py @@ -424,6 +424,16 @@ def test_particle_slice(self): with self.assertRaises(TypeError): system.part[[ids[0], 1.2]] + def test_to_dict(self): + self.system.part.clear() + p = self.system.part.add( + pos=np.random.uniform(size=(10, 3)) * self.system.box_l) + pp = str(p) + pdict = p.to_dict() + p.remove() + self.system.part.add(pdict) + self.assertEqual(str(self.system.part.select()), pp) + if __name__ == "__main__": ut.main()