Save beam spectra to external pickle file

- On serializing, create an external pickle file with the beam spectra and replace the spectra arrays in the dictionary with indices into the pickle file
- On deserializing, read the external pickle file back and set the spectra appropiately
- Change to_dict and from_dict to use flex.doubles instead of lists of floats. Not a primitive dictionary any longer but much faster to read/write

model/__init__.py

@@ -796,6 +796,8 @@ def as_json(self, filename=None, compact=False, split=False):
         from dxtbx.datablock import AutoEncoder
 
         for fname, obj in to_write:
+            beam_spectra_to_pickle(filename, obj)
+
             if compact:
                 separators = (",", ":")
                 indent = None
@@ -856,3 +858,47 @@ def from_file(filename, check_format=True):
         from .experiment_list import ExperimentListFactory
 
         return ExperimentListFactory.from_serialized_format(filename, check_format)
+
+
+def beam_spectra_to_pickle(filename, obj):
+    """
+    Helper function to take an experiment list dictionary and move the
+    spectra to a pickle file
+
+    Args:
+        filename: File name the experiment list will be saved to
+        obj: Experiment list dictionary
+    """
+    if "beam" not in obj:
+        return
+    all_energies = []
+    all_weights = []
+    test = None
+    all_eq = True
+    for beam in obj["beam"]:
+        if "spectrum_energies" not in beam:
+            continue
+        if test is None:
+            test = beam["spectrum_energies"]
+        else:
+            if (test == beam["spectrum_energies"]).count(False):
+                all_eq = False
+                break
+
+    for beam in obj["beam"]:
+        if "spectrum_energies" not in beam:
+            continue
+        if not all_eq or not all_energies:
+            all_energies.append(beam["spectrum_energies"])
+        all_weights.append(beam["spectrum_weights"])
+        del beam["spectrum_energies"]
+        del beam["spectrum_weights"]
+        beam["spectrum_index"] = len(all_weights) - 1
+    if not all_energies:
+        return
+    pickle_filename = os.path.splitext(filename)[0] + "_spectrum.pickle"
+    obj["spectra_pickle"] = pickle_filename
+    with open(str(pickle_filename), "wb") as outfile:
+        pickle.dump(
+            (all_energies, all_weights), outfile, protocol=pickle.HIGHEST_PROTOCOL
+        )

model/boost_python/beam.cc

@@ -215,14 +215,8 @@ namespace dxtbx { namespace model { namespace boost_python {
       result["s0_at_scan_points"] = l;
     }
     if (obj.get_spectrum_energies().size()) {
-      boost::python::list e;
-      boost::python::list w;
-      for (size_t i = 0; i < obj.get_spectrum_energies().size(); i++) {
-        e.append(obj.get_spectrum_energies()[i]);
-        w.append(obj.get_spectrum_weights()[i]);
-      }
-      result["spectrum_energies"] = e;
-      result["spectrum_weights"] = w;
+      result["spectrum_energies"] = obj.get_spectrum_energies();
+      result["spectrum_weights"] = obj.get_spectrum_weights();
     }
     return result;
   }
@@ -246,16 +240,8 @@ namespace dxtbx { namespace model { namespace boost_python {
     }
     if (obj.has_key("spectrum_energies")) {
       DXTBX_ASSERT(obj.has_key("spectrum_weights"));
-      boost::python::list e_list = boost::python::extract<boost::python::list>(obj["spectrum_energies"]);
-      boost::python::list w_list = boost::python::extract<boost::python::list>(obj["spectrum_weights"]);
-      DXTBX_ASSERT(boost::python::len(e_list) == boost::python::len(w_list));
-      scitbx::af::shared<double> e;
-      scitbx::af::shared<double> w;
-      for (size_t i = 0; i < boost::python::len(e_list); i++) {
-        e.push_back(boost::python::extract<double>(e_list[i]));
-        w.push_back(boost::python::extract<double>(w_list[i]));
-      }
-      b->set_spectrum(e, w);
+      b->set_spectrum(boost::python::extract<scitbx::af::shared<double> >(obj["spectrum_energies"]),
+                      boost::python::extract<scitbx::af::shared<double> >(obj["spectrum_energies"]));
     }
     return b;
   }

model/experiment_list.py

@@ -371,9 +371,29 @@ def decode(self):
                 )
             )
 
+        # Read the beam spectra if available
+        self._read_spectra_pickle(el.beams())
+
         # Return the experiment list
         return el
 
+    def _read_spectra_pickle(self, beams):
+        """ Read the beam spectra if available """
+        if "spectra_pickle" not in self._obj:
+            return
+        with open(self._obj["spectra_pickle"], "rb") as infile:
+            all_energies, all_weights = pickle.load(infile)
+        assert len(all_energies) == 1 or len(all_energies) == len(all_weights)
+        for i, beam in enumerate(beams):
+            spectrum_index = self._obj["beam"][i].get("spectrum_index")
+            if spectrum_index is not None:
+                if len(all_energies) == 1:
+                    beam.set_spectrum(all_energies[0], all_weights[spectrum_index])
+                else:
+                    beam.set_spectrum(
+                        all_energies[spectrum_index], all_weights[spectrum_index]
+                    )
+
     def _make_mem_imageset(self, imageset):
         """Can't make a mem imageset from dict."""
         return None