Add __reduce__() to box variants to allow pickling

glotzerlab · Mar 12, 2024 · 6c17eab · 6c17eab
1 parent f66b49f
commit 6c17eab
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/hoomd/variant/box.py b/hoomd/variant/box.py
@@ -128,6 +128,10 @@ def __init__(self, initial_box, final_box, variant):
         _hoomd.VectorVariantBoxInterpolate.__init__(self, box1._cpp_obj,
                                                     box2._cpp_obj, variant)
 
+    def __reduce__(self):
+        """Reduce values to picklable format."""
+        return (type(self), (self.initial_box, self.final_box, self.variant))
+
     @property
     def initial_box(self):
         """hoomd.Box: The initial box."""
@@ -191,6 +195,11 @@ def __init__(self, initial_box, final_volume, t_start, t_ramp):
         _hoomd.VectorVariantBoxInverseVolumeRamp.__init__(
             self, box._cpp_obj, final_volume, t_start, t_ramp)
 
+    def __reduce__(self):
+        """Reduce values to picklable format."""
+        return (type(self), (self.initial_box, self.final_volume, self.t_start,
+                             self.t_ramp))
+
     @property
     def initial_box(self):
         """hoomd.Box: The initial box."""