From 1baf8c6e056b89810e78922239e885f59e936a6e Mon Sep 17 00:00:00 2001
From: Ben Webb <ben@salilab.org>
Date: Wed, 25 Oct 2023 16:03:40 -0700
Subject: [PATCH] Serialize DirectionMover

---
 modules/core/include/DirectionMover.h     | 15 +++++++++++-
 modules/core/pyext/swig.i-in              |  2 +-
 modules/core/src/DirectionMover.cpp       |  2 ++
 modules/core/test/test_direction_mover.py | 30 +++++++++++++++++++++++
 4 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/modules/core/include/DirectionMover.h b/modules/core/include/DirectionMover.h
index a1fcc365b5..856b5450f9 100644
--- a/modules/core/include/DirectionMover.h
+++ b/modules/core/include/DirectionMover.h
@@ -2,7 +2,7 @@
  *  \file IMP/core/DirectionMover.h
  *  \brief A mover that transforms a Direction.
  *
- *  Copyright 2007-2022 IMP Inventors. All rights reserved.
+ *  Copyright 2007-2023 IMP Inventors. All rights reserved.
  *
  */
 
@@ -16,6 +16,9 @@
 #include <IMP/Particle.h>
 #include <IMP/Object.h>
 #include <IMP/Model.h>
+#include <cereal/access.hpp>
+#include <cereal/types/base_class.hpp>
+#include <cereal/types/polymorphic.hpp>
 
 
 IMPCORE_BEGIN_NAMESPACE
@@ -36,12 +39,22 @@ class IMPCOREEXPORT DirectionMover : public MonteCarloMover {
     void initialize(ParticleIndex pi, double max_rotation,
                     double reflect_probability);
 
+    friend class cereal::access;
+
+    template<class Archive> void serialize(Archive &ar) {
+      ar(cereal::base_class<MonteCarloMover>(this), last_direction_,
+         max_angle_, reflect_prob_, pi_);
+    }
+    IMP_OBJECT_SERIALIZE_DECL(DirectionMover);
+
   public:
     DirectionMover(Model *m, ParticleIndex pi, Float max_rotation,
                    Float reflect_probability);
 
     DirectionMover(Direction d, Float max_rotation,
                    Float reflect_probability);
+
+    DirectionMover() {}
   
     //! Set the maximum rotation in radians.
     void set_maximum_rotation(Float mr);
diff --git a/modules/core/pyext/swig.i-in b/modules/core/pyext/swig.i-in
index 36b88dada7..faf3336f5f 100644
--- a/modules/core/pyext/swig.i-in
+++ b/modules/core/pyext/swig.i-in
@@ -12,7 +12,7 @@ IMP_SWIG_OBJECT_INSTANCE( IMP::core, AttributeSingletonScore, AttributeSingleton
 IMP_SWIG_OBJECT_SERIALIZE( IMP::core, BallMover, BallMovers);
 IMP_SWIG_OBJECT_SERIALIZE( IMP::core, SerialMover, SerialMovers);
 IMP_SWIG_OBJECT_SERIALIZE( IMP::core, SubsetMover, SubsetMovers);
-IMP_SWIG_OBJECT( IMP::core, DirectionMover, DirectionMovers);
+IMP_SWIG_OBJECT_SERIALIZE( IMP::core, DirectionMover, DirectionMovers);
 IMP_SWIG_OBJECT( IMP::core, SurfaceMover, SurfaceMovers);
 IMP_SWIG_OBJECT_INSTANCE( IMP::core, BoundingBox3DSingletonScore, BoundingBox3DSingletonScore, BoundingBox3DSingletonScores);
 IMP_SWIG_OBJECT_INSTANCE( IMP::core, BoundingSphere3DSingletonScore, BoundingSphere3DSingletonScore, BoundingSphere3DSingletonScores);
diff --git a/modules/core/src/DirectionMover.cpp b/modules/core/src/DirectionMover.cpp
index bedefb698b..63a074adfc 100644
--- a/modules/core/src/DirectionMover.cpp
+++ b/modules/core/src/DirectionMover.cpp
@@ -81,4 +81,6 @@ ModelObjectsTemp DirectionMover::do_get_inputs() const {
   return ParticlesTemp(1, get_model()->get_particle(pi_));
 }
 
+IMP_OBJECT_SERIALIZE_IMPL(IMP::core::DirectionMover);
+
 IMPCORE_END_NAMESPACE
diff --git a/modules/core/test/test_direction_mover.py b/modules/core/test/test_direction_mover.py
index c8df2b2869..fe4c5ff4bf 100644
--- a/modules/core/test/test_direction_mover.py
+++ b/modules/core/test/test_direction_mover.py
@@ -3,6 +3,7 @@
 import IMP.core
 import IMP.atom
 import IMP.test
+import pickle
 
 
 class Tests(IMP.test.TestCase):
@@ -62,6 +63,35 @@ def test_inputs(self):
         self.assertSetEqual(set([d.get_particle()]), set(mv.get_inputs()))
         mv.set_was_used(True)
 
+    def test_pickle(self):
+        """Test (un-)pickle of DirectionMover"""
+        m = IMP.Model()
+        d = IMP.core.Direction.setup_particle(IMP.Particle(m),
+                                              IMP.algebra.Vector3D(0, 0, 1))
+        d.set_direction_is_optimized(True)
+        mvr = IMP.core.DirectionMover(d, .1, 1.)
+        mvr.set_name("foo")
+        dump = pickle.dumps(mvr)
+
+        newmvr = pickle.loads(dump)
+        self.assertEqual(newmvr.get_name(), "foo")
+        self.assertSetEqual(set([d.get_particle()]), set(newmvr.get_inputs()))
+
+    def test_pickle_polymorphic(self):
+        """Test (un-)pickle of DirectionMover via polymorphic pointer"""
+        m = IMP.Model()
+        d = IMP.core.Direction.setup_particle(IMP.Particle(m),
+                                              IMP.algebra.Vector3D(0, 0, 1))
+        d.set_direction_is_optimized(True)
+        mvr = IMP.core.DirectionMover(d, .1, 1.)
+        mvr.set_name("foo")
+        sm = IMP.core.SerialMover([mvr])
+        dump = pickle.dumps(sm)
+
+        newsm = pickle.loads(dump)
+        newmvr, = newsm.get_movers()
+        self.assertEqual(newmvr.get_name(), "foo")
+
 
 if __name__ == '__main__':
     IMP.test.main()