Serialize vector metrics

salilab · Oct 24, 2023 · 85f5269 · 85f5269
1 parent 6c5868c
commit 85f5269
Show file tree

Hide file tree

Showing 4 changed files with 72 additions and 6 deletions.
diff --git a/modules/algebra/include/vector_metrics.h b/modules/algebra/include/vector_metrics.h
@@ -1,7 +1,7 @@
 /**
  *  \file IMP/algebra/vector_metrics.h   \brief Functions to generate vectors.
  *
- *  Copyright 2007-2022 IMP Inventors. All rights reserved.
+ *  Copyright 2007-2023 IMP Inventors. All rights reserved.
  *
  */
 
@@ -17,6 +17,9 @@
 #include <IMP/WeakPointer.h>
 #include <algorithm>
 #include <numeric>
+#include <cereal/access.hpp>
+#include <cereal/types/polymorphic.hpp>
+#include <cereal/types/base_class.hpp>
 
 IMPALGEBRA_BEGIN_NAMESPACE
 
@@ -32,6 +35,11 @@ IMP_OBJECTS(VectorKDMetric, VectorKDMetrics);
 
 //! The l2 norm on the distance vector.
 class IMPALGEBRAEXPORT EuclideanVectorKDMetric : public VectorKDMetric {
+  friend class cereal::access;
+  template<class Archive> void serialize(Archive &ar) {
+    ar(cereal::base_class<Object>(this));
+  }
+  IMP_OBJECT_SERIALIZE_DECL(EuclideanVectorKDMetric);
  public:
   EuclideanVectorKDMetric(std::string name = "EuclideanVectorKDMetric%1%");
   double get_distance(const VectorKD &a, const VectorKD &b) const override {
@@ -50,6 +58,11 @@ class IMPALGEBRAEXPORT EuclideanVectorKDMetric : public VectorKDMetric {
     centroid is the center of the bounding box of the vectors.
  */
 class IMPALGEBRAEXPORT MaxVectorKDMetric : public VectorKDMetric {
+  friend class cereal::access;
+  template<class Archive> void serialize(Archive &ar) {
+    ar(cereal::base_class<Object>(this));
+  }
+  IMP_OBJECT_SERIALIZE_DECL(MaxVectorKDMetric);
  public:
   MaxVectorKDMetric(std::string name = "MaxVectorKDMetric%1%");
   double get_distance(const VectorKD &a, const VectorKD &b) const override {

diff --git a/modules/algebra/pyext/swig.i-in b/modules/algebra/pyext/swig.i-in
@@ -26,8 +26,8 @@ IMP_SWIG_EIGEN_VECTOR(VectorXf);
 IMP_SWIG_EIGEN_VECTOR(Vector3d);
 IMP_SWIG_EIGEN_MATRIX(ArrayXXf);
 IMP_SWIG_EIGEN_MATRIX(ArrayXXd);
-IMP_SWIG_OBJECT(IMP::algebra, EuclideanVectorKDMetric, EuclideanVectorKDMetrics);
-IMP_SWIG_OBJECT(IMP::algebra, MaxVectorKDMetric, MaxVectorKDMetrics);
+IMP_SWIG_OBJECT_SERIALIZE(IMP::algebra, EuclideanVectorKDMetric, EuclideanVectorKDMetrics);
+IMP_SWIG_OBJECT_SERIALIZE(IMP::algebra, MaxVectorKDMetric, MaxVectorKDMetrics);
 IMP_SWIG_OBJECT(IMP::algebra, DynamicNearestNeighbor3D, DynamicNearestNeighbor3Ds);
 
 IMP_SWIG_VALUE_SERIALIZE(IMP::algebra, Rotation2D, Rotation2Ds);

diff --git a/modules/algebra/src/vector_metrics.cpp b/modules/algebra/src/vector_metrics.cpp
@@ -1,8 +1,7 @@
 /**
- *  \file vector_generators.cpp
- *  \brief Support for rigid bodies.
+ *  \file vector_metrics.cpp   \brief Functions to generate vectors.
  *
- *  Copyright 2007-2022 IMP Inventors. All rights reserved.
+ *  Copyright 2007-2023 IMP Inventors. All rights reserved.
  *
  */
 
@@ -14,5 +13,10 @@ VectorKDMetric::VectorKDMetric(std::string name) : Object(name) {}
 
 EuclideanVectorKDMetric::EuclideanVectorKDMetric(std::string name)
     : VectorKDMetric(name) {}
+
 MaxVectorKDMetric::MaxVectorKDMetric(std::string name) : VectorKDMetric(name) {}
+
+IMP_OBJECT_SERIALIZE_IMPL(IMP::algebra::EuclideanVectorKDMetric);
+IMP_OBJECT_SERIALIZE_IMPL(IMP::algebra::MaxVectorKDMetric);
+
 IMPALGEBRA_END_NAMESPACE
diff --git a/modules/algebra/test/test_vector_metric.py b/modules/algebra/test/test_vector_metric.py
@@ -1,6 +1,7 @@
 import IMP
 import IMP.test
 import IMP.algebra
+import pickle
 
 
 class Tests(IMP.test.TestCase):
@@ -14,6 +15,30 @@ def test_euc_vector_metric(self):
         self.assertLess(IMP.algebra.get_distance(
             m.get_centroid(vs), V(5.33333, 6.33333, 7.33333)), 1e-4)
 
+    def test_euc_vector_pickle(self):
+        """Check (un-)pickle of Euclidean vector metric"""
+        m = IMP.algebra.EuclideanVectorKDMetric("foo")
+        dump = pickle.dumps(m)
+        newm = pickle.loads(dump)
+        self.assertEqual(newm.get_name(), "foo")
+        V = IMP.algebra.VectorKD
+        vs = [V(1, 2, 3), V(5, 6, 7), V(10, 11, 12)]
+        self.assertAlmostEqual(newm.get_distance(vs[0], vs[1]),
+                               6.928, delta=1e-3)
+
+    def test_euc_vector_pickle_polymorphic(self):
+        """Check (un-)pickle of Euclidean vector metric via polymorphic ptr"""
+        m = IMP.Model()
+        mk = IMP.ModelKey("mk")
+        metric = IMP.algebra.EuclideanVectorKDMetric("foo")
+        m.add_data(mk, metric)
+        dump = pickle.dumps(m)
+        newm = pickle.loads(dump)
+
+        newmetric = IMP.algebra.EuclideanVectorKDMetric.get_from(
+            newm.get_data(mk))
+        self.assertEqual(newmetric.get_name(), "foo")
+
     def test_max_vector_metric(self):
         """Check max vector metric"""
         m = IMP.algebra.MaxVectorKDMetric("foo")
@@ -23,6 +48,30 @@ def test_max_vector_metric(self):
         self.assertLess(IMP.algebra.get_distance(m.get_centroid(vs),
                                                  V(5.5, 6.5, 7.5)), 1e-4)
 
+    def test_max_vector_metric_pickle(self):
+        """Check (un-)pickle of max vector metric"""
+        m = IMP.algebra.MaxVectorKDMetric("foo")
+        dump = pickle.dumps(m)
+        newm = pickle.loads(dump)
+        self.assertEqual(newm.get_name(), "foo")
+        V = IMP.algebra.VectorKD
+        vs = [V(1, 2, 3), V(5, 6, 7), V(10, 11, 12)]
+        self.assertAlmostEqual(newm.get_distance(vs[0], vs[1]),
+                               4.0, delta=1e-3)
+
+    def test_max_vector_pickle_polymorphic(self):
+        """Check (un-)pickle of max vector metric via polymorphic ptr"""
+        m = IMP.Model()
+        mk = IMP.ModelKey("mk")
+        metric = IMP.algebra.MaxVectorKDMetric("foo")
+        m.add_data(mk, metric)
+        dump = pickle.dumps(m)
+        newm = pickle.loads(dump)
+
+        newmetric = IMP.algebra.MaxVectorKDMetric.get_from(
+            newm.get_data(mk))
+        self.assertEqual(newmetric.get_name(), "foo")
+
 
 if __name__ == '__main__':
     IMP.test.main()