Add spectrum to beam object

model/__init__.py

@@ -798,6 +798,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_h5(filename, obj)
+
             if compact:
                 separators = (",", ":")
                 indent = None
@@ -887,3 +889,64 @@ def change_basis(self, change_of_basis_ops, in_place=False):
         for expt, cb_op in zip(return_expts, change_of_basis_ops):
             expt.crystal = expt.crystal.change_basis(cb_op)
         return return_expts
+
+def beam_spectra_to_h5(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
+
+    import h5py
+
+    h5_filename = os.path.splitext(filename)[0] + "_spectrum.h5"
+    obj["spectra_h5"] = h5_filename
+    f = h5py.File(h5_filename, "w")
+
+    def get_stacked_data(data):
+        if len(data) > 1:
+            import numpy as np
+
+            return np.vstack([array.as_numpy_array() for array in data])
+        else:
+            return data[0].as_numpy_array()
+
+    all_energies = get_stacked_data(all_energies)
+    all_weights = get_stacked_data(all_weights)
+
+    f.create_dataset(
+        "all_energies", all_energies.shape, data=all_energies, dtype=all_energies.dtype
+    )
+    f.create_dataset(
+        "all_weights", all_weights.shape, data=all_weights, dtype=all_weights.dtype
+    )

model/beam.h

@@ -80,14 +80,23 @@ namespace dxtbx { namespace model {
     virtual vec3<double> get_s0_at_scan_point(std::size_t index) const = 0;
     // Reset the scan points
     virtual void reset_scan_points() = 0;
+    // Set the spectrum
+    virtual void set_spectrum(scitbx::af::shared<double> spectrum_energies, scitbx::af::shared<double>) = 0;
+    // Get the spectrum energies
+    virtual scitbx::af::shared<double> get_spectrum_energies() const = 0;
+    // Get the spectrum weights
+    virtual scitbx::af::shared<double> get_spectrum_weights() const = 0;
+    // Get the spectrum weighted wavelength
+    virtual double get_weighted_wavelength() const = 0;
     // Check wavelength and direction are (almost) same
     virtual bool operator==(const BeamBase &rhs) const = 0;
     // Check if two models are similar
     virtual bool is_similar_to(const BeamBase &rhs,
                                double wavelength_tolerance,
                                double direction_tolerance,
                                double polarization_normal_tolerance,
-                               double polarization_fraction_tolerance) const = 0;
+                               double polarization_fraction_tolerance,
+                               double spectrum_weighted_wavelength_tolerance) const = 0;
     // Check wavelength and direction are not (almost) equal.
     virtual bool operator!=(const BeamBase &rhs) const = 0;
     //  Rotate the beam about an axis
@@ -392,7 +401,8 @@ namespace dxtbx { namespace model {
                        double wavelength_tolerance,
                        double direction_tolerance,
                        double polarization_normal_tolerance,
-                       double polarization_fraction_tolerance) const {
+                       double polarization_fraction_tolerance,
+                       double spectrum_weighted_wavelength_tolerance) const {
       // scan varying model checks
       if (get_num_scan_points() != rhs.get_num_scan_points()) {
         return false;
@@ -414,6 +424,14 @@ namespace dxtbx { namespace model {
         }
       }
 
+      // spectrum checks
+      if (spectrum_energies_.size() != rhs.get_spectrum_energies().size())
+        return false;
+      if (spectrum_energies_.size())
+        if (!(std::abs(get_weighted_wavelength() - rhs.get_weighted_wavelength()) <=
+                      spectrum_weighted_wavelength_tolerance))
+          return false;
+
       // static model checks
       return std::abs(angle_safe(direction_, rhs.get_sample_to_source_direction()))
                <= direction_tolerance
@@ -439,6 +457,42 @@ namespace dxtbx { namespace model {
     }
 
     friend std::ostream &operator<<(std::ostream &os, const Beam &b);
+    /**
+     * Set the spectrum. X is eV, Y is dimensionless
+     */
+    void set_spectrum(scitbx::af::shared<double> spectrum_energies, scitbx::af::shared<double> spectrum_weights) {
+      DXTBX_ASSERT(spectrum_energies.size() == spectrum_weights.size());
+      spectrum_energies_ = spectrum_energies;
+      spectrum_weights_ = spectrum_weights;
+    }
+
+    /**
+     * Get the spectrum energies (eV)
+     */
+    scitbx::af::shared<double> get_spectrum_energies() const {
+      return spectrum_energies_;
+    }
+
+    /**
+     * Get the spectrum weights
+     */
+    scitbx::af::shared<double> get_spectrum_weights() const {
+      return spectrum_weights_;
+    }
+
+    /**
+     * Get the spectrum weighted wavelength
+     */
+    double get_weighted_wavelength() const {
+      double weighted_sum = 0;
+      double summed_weights = 0;
+      for (size_t i = 0; i < spectrum_energies_.size(); i++) {
+        weighted_sum += spectrum_energies_[i] * spectrum_weights_[i];
+        summed_weights += spectrum_weights_[i];
+      }
+      DXTBX_ASSERT(weighted_sum > 0 && summed_weights > 0);
+      return 12398.4187 / (weighted_sum / summed_weights); //  eV per Å conversion factor
+    }
 
   private:
     double wavelength_;
@@ -450,6 +504,8 @@ namespace dxtbx { namespace model {
     double flux_;
     double transmission_;
     scitbx::af::shared<vec3<double> > s0_at_scan_points_;
+    scitbx::af::shared<double> spectrum_energies_;
+    scitbx::af::shared<double> spectrum_weights_;
   };
 
   /** Print beam information */
@@ -463,6 +519,8 @@ namespace dxtbx { namespace model {
     os << "    polarization normal: " << b.get_polarization_normal().const_ref()
        << "\n";
     os << "    polarization fraction: " << b.get_polarization_fraction() << "\n";
+    if (b.get_spectrum_energies().size())
+      os << "    spectrum weighted wavelength: " << b.get_weighted_wavelength() << "\n";
     return os;
   }
 

model/boost_python/beam.cc

@@ -38,7 +38,9 @@ namespace dxtbx { namespace model { namespace boost_python {
                                        obj.get_polarization_normal(),
                                        obj.get_polarization_fraction(),
                                        obj.get_flux(),
-                                       obj.get_transmission());
+                                       obj.get_transmission(),
+                                       obj.get_spectrum_energies(),
+                                       obj.get_spectrum_weights());
     }
 
     static boost::python::tuple getstate(boost::python::object obj) {
@@ -97,7 +99,7 @@ namespace dxtbx { namespace model { namespace boost_python {
     return beam;
   }
 
-  static Beam *make_beam_w_all(vec3<double> sample_to_source,
+  static Beam *make_beam_w_most(vec3<double> sample_to_source,
                                double wavelength,
                                double divergence,
                                double sigma_divergence,
@@ -129,6 +131,23 @@ namespace dxtbx { namespace model { namespace boost_python {
     return beam;
   }
 
+  static Beam *make_beam_w_all(vec3<double> sample_to_source,
+                               double wavelength,
+                               double divergence,
+                               double sigma_divergence,
+                               vec3<double> polarization_normal,
+                               double polarization_fraction,
+                               double flux,
+                               double transmission,
+                               scitbx::af::shared<double> spectrum_energies,
+                               scitbx::af::shared<double> spectrum_weights,
+                               bool deg) {
+    Beam *beam = make_beam_w_most(sample_to_source, wavelength, divergence, sigma_divergence,
+                                  polarization_normal, polarization_fraction, flux, transmission, deg);
+    beam->set_spectrum(spectrum_energies, spectrum_weights);
+    return beam;
+  }
+
   static double get_divergence(const Beam &beam, bool deg) {
     double divergence = beam.get_divergence();
     return deg ? rad_as_deg(divergence) : divergence;
@@ -195,6 +214,10 @@ namespace dxtbx { namespace model { namespace boost_python {
       }
       result["s0_at_scan_points"] = l;
     }
+    if (obj.get_spectrum_energies().size()) {
+      result["spectrum_energies"] = obj.get_spectrum_energies();
+      result["spectrum_weights"] = obj.get_spectrum_weights();
+    }
     return result;
   }
 
@@ -215,6 +238,11 @@ namespace dxtbx { namespace model { namespace boost_python {
         boost::python::extract<boost::python::list>(obj["s0_at_scan_points"]);
       Beam_set_s0_at_scan_points_from_list(*b, s0_at_scan_points);
     }
+    if (obj.has_key("spectrum_energies")) {
+      DXTBX_ASSERT(obj.has_key("spectrum_weights"));
+      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;
   }
 
@@ -262,7 +290,8 @@ namespace dxtbx { namespace model { namespace boost_python {
             arg("wavelength_tolerance") = 1e-6,
             arg("direction_tolerance") = 1e-6,
             arg("polarization_normal_tolerance") = 1e-6,
-            arg("polarization_fraction_tolerance") = 1e-6));
+            arg("polarization_fraction_tolerance") = 1e-6,
+            arg("spectrum_weighted_wavelength_tolerance") = 1e-6));
 
     // Export Beam : BeamBase
     class_<Beam, boost::shared_ptr<Beam>, bases<BeamBase> >("Beam")
@@ -284,6 +313,18 @@ namespace dxtbx { namespace model { namespace boost_python {
           default_call_policies(),
           (arg("s0"), arg("divergence"), arg("sigma_divergence"), arg("deg") = true)))
       .def("__init__",
+           make_constructor(&make_beam_w_most,
+                            default_call_policies(),
+                            (arg("direction"),
+                             arg("wavelength"),
+                             arg("divergence"),
+                             arg("sigma_divergence"),
+                             arg("polarization_normal"),
+                             arg("polarization_fraction"),
+                             arg("flux"),
+                             arg("transmission"),
+                             arg("deg") = true)))
+       .def("__init__",
            make_constructor(&make_beam_w_all,
                             default_call_policies(),
                             (arg("direction"),
@@ -294,8 +335,18 @@ namespace dxtbx { namespace model { namespace boost_python {
                              arg("polarization_fraction"),
                              arg("flux"),
                              arg("transmission"),
+                             arg("spectrum_energies"),
+                             arg("spectrum_weights"),
                              arg("deg") = true)))
       .def("__str__", &beam_to_string)
+      .def("get_spectrum_energies",
+          &Beam::get_spectrum_energies)
+      .def("get_spectrum_weights",
+          &Beam::get_spectrum_weights)
+      .def("get_weighted_wavelength",
+          &Beam::get_weighted_wavelength)
+      .def("set_spectrum",
+          &Beam::set_spectrum)
       .def("to_dict", &to_dict<Beam>)
       .def("from_dict", &from_dict<Beam>, return_value_policy<manage_new_object>())
       .staticmethod("from_dict")

model/experiment_list.py

@@ -372,9 +372,41 @@ def decode(self):
                 )
             )
 
+        # Read the beam spectra if available
+        self._read_spectra_h5(el.beams())
+
         # Return the experiment list
         return el
 
+    def _read_spectra_h5(self, beams):
+        """ Read the beam spectra if available """
+        if "spectra_h5" not in self._obj:
+            return
+
+        import h5py
+
+        f = h5py.File(os.path.join(self._directory, self._obj["spectra_h5"]), "r")
+        all_energies = f["all_energies"][()]
+        all_weights = f["all_weights"][()]
+
+        assert (
+            len(all_energies.shape) == 1 or all_energies.shape == all_weights.shape
+        ), (all_energies.shape, all_weights.shape)
+        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.shape) == 1:
+                    spectrum = all_energies
+                else:
+                    spectrum = all_energies[spectrum_index]
+
+                if len(all_weights.shape) == 1:
+                    weights = all_weights
+                else:
+                    weights = all_weights[spectrum_index]
+
+                beam.set_spectrum(spectrum, weights)
+
     def _make_mem_imageset(self, imageset):
         """Can't make a mem imageset from dict."""
         return None

tests/model/test_beam.py

@@ -6,6 +6,7 @@
 
 from libtbx.phil import parse
 from scitbx import matrix
+from scitbx.array_family import flex
 
 from dxtbx.model import Beam
 from dxtbx.model.beam import BeamFactory, beam_phil_scope
@@ -178,3 +179,22 @@ def test_beam_object_comparison():
 def test_beam_self_serialization():
     beam = Beam()
     assert beam == BeamFactory.from_dict(beam.to_dict())
+
+
+def test_spectrum_beam():
+    spectrum_energies = flex.double(range(9450, 9550))
+    spectrum_weights = flex.double(range(len(spectrum_energies)))
+    b1 = Beam()
+    b2 = Beam()
+    b1.set_spectrum(spectrum_energies, spectrum_weights)
+    b2.set_spectrum(spectrum_energies, spectrum_weights)
+    assert b1.get_weighted_wavelength() == pytest.approx(1.3028567060142213)
+    assert b1.is_similar_to(b2)
+    b2.set_spectrum(spectrum_energies + 50, spectrum_weights)
+    assert not b1.is_similar_to(b2)
+
+    b3 = Beam()
+    b1.set_wavelength(1.2)
+    b3.set_wavelength(1.2)
+    assert not b1.is_similar_to(b3)
+    assert not b3.is_similar_to(b1)