Implement __getstate__ and __setstate__

llr-cta · Oct 25, 2023 · 2817d35 · 2817d35
1 parent 3b4eab1
commit 2817d35
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 13 deletions.
diff --git a/include/math/nspace.hpp b/include/math/nspace.hpp
@@ -47,6 +47,7 @@ class TreeSparseNSpace
   TreeSparseNSpace(const Eigen::VectorXd& xlo, const Eigen::VectorXd& xhi,
     const Eigen::VectorXi& n);
   TreeSparseNSpace(const std::vector<Axis>& axes);
+  TreeSparseNSpace(const calin::ix::math::nspace::NSpaceData& proto);
 
   void clear() { bins_.clear(); }
   void injest(const TreeSparseNSpace& o);
@@ -210,6 +211,7 @@ class BlockSparseNSpace
   BlockSparseNSpace(const Eigen::VectorXd& xlo, const Eigen::VectorXd& xhi,
     const Eigen::VectorXi& n, unsigned log2_block_size = 0);
   BlockSparseNSpace(const std::vector<Axis>& axes, unsigned log2_block_size = 0);
+  BlockSparseNSpace(const calin::ix::math::nspace::NSpaceData& proto, unsigned log2_block_size = 0);
 
   ~BlockSparseNSpace();
 

diff --git a/src/math/nspace.cpp b/src/math/nspace.cpp
@@ -84,6 +84,15 @@ TreeSparseNSpace::TreeSparseNSpace(const std::vector<Axis>& axes):
   // nothing to see here
 }
 
+TreeSparseNSpace::TreeSparseNSpace(const calin::ix::math::nspace::NSpaceData& proto):
+  TreeSparseNSpace(
+    Eigen::Map<const Eigen::VectorXd>(proto.axes_lower_bounds().data(), proto.axes_lower_bounds().size()),
+    Eigen::Map<const Eigen::VectorXd>(proto.axes_upper_bounds().data(), proto.axes_upper_bounds().size()),
+    Eigen::Map<const Eigen::VectorXi>(proto.axes_num_bins().data(), proto.axes_num_bins().size()))
+{
+  this->accumulate_from_proto(proto);
+}
+
 void TreeSparseNSpace::injest(const TreeSparseNSpace& o)
 {
   if(xlo_!=o.xlo_ or xhi_!=o.xhi_ or n_!=o.n_) {
@@ -407,12 +416,7 @@ void TreeSparseNSpace::accumulate_from_proto(const calin::ix::math::nspace::NSpa
 
 TreeSparseNSpace* TreeSparseNSpace::create_from_proto(const calin::ix::math::nspace::NSpaceData& proto)
 {
-  TreeSparseNSpace* nspace = new TreeSparseNSpace(
-    Eigen::Map<const Eigen::VectorXd>(proto.axes_lower_bounds().data(), proto.axes_lower_bounds().size()),
-    Eigen::Map<const Eigen::VectorXd>(proto.axes_upper_bounds().data(), proto.axes_upper_bounds().size()),
-    Eigen::Map<const Eigen::VectorXi>(proto.axes_num_bins().data(), proto.axes_num_bins().size()));
-  nspace->accumulate_from_proto(proto);
-  return nspace;
+  return new TreeSparseNSpace(proto);
 }
 
 // =============================================================================
@@ -467,6 +471,16 @@ BlockSparseNSpace::BlockSparseNSpace(const std::vector<Axis>& axes, unsigned log
   // nothing to see here
 }
 
+BlockSparseNSpace::BlockSparseNSpace(const calin::ix::math::nspace::NSpaceData& proto, unsigned log2_block_size ):
+  BlockSparseNSpace(
+    Eigen::Map<const Eigen::VectorXd>(proto.axes_lower_bounds().data(), proto.axes_lower_bounds().size()),
+    Eigen::Map<const Eigen::VectorXd>(proto.axes_upper_bounds().data(), proto.axes_upper_bounds().size()),
+    Eigen::Map<const Eigen::VectorXi>(proto.axes_num_bins().data(), proto.axes_num_bins().size()),
+    log2_block_size)
+{
+  this->accumulate_from_proto(proto);
+}
+
 BlockSparseNSpace::~BlockSparseNSpace()
 {
   for(auto* array : alloc_all_list_) {
@@ -1272,13 +1286,7 @@ void BlockSparseNSpace::accumulate_from_proto(const calin::ix::math::nspace::NSp
 BlockSparseNSpace* BlockSparseNSpace::create_from_proto(const calin::ix::math::nspace::NSpaceData& proto, 
   unsigned log2_block_size)
 {
-  BlockSparseNSpace* nspace = new BlockSparseNSpace(
-    Eigen::Map<const Eigen::VectorXd>(proto.axes_lower_bounds().data(), proto.axes_lower_bounds().size()),
-    Eigen::Map<const Eigen::VectorXd>(proto.axes_upper_bounds().data(), proto.axes_upper_bounds().size()),
-    Eigen::Map<const Eigen::VectorXi>(proto.axes_num_bins().data(), proto.axes_num_bins().size()),
-    log2_block_size);
-  nspace->accumulate_from_proto(proto);
-  return nspace;
+  return new BlockSparseNSpace(proto, log2_block_size);
 }
 
 #if 0

diff --git a/swig/math/nspace.i b/swig/math/nspace.i
@@ -48,4 +48,25 @@
 %newobject create_from_proto;
 %newobject as_proto;
 
+%extend calin::math::nspace::TreeSparseNSpace {
+  %pythoncode %{
+    def __getstate__(self):
+      return self.as_proto()
+
+    def __setstate__(self, proto):
+      self.__init__(proto)
+  %}
+}
+
+
+%extend calin::math::nspace::BlockSparseNSpace {
+  %pythoncode %{
+    def __getstate__(self):
+      return self.as_proto()
+
+    def __setstate__(self, proto):
+      self.__init__(proto)
+  %}
+}
+
 %include "math/nspace.hpp"