Add Spectrum as a read-only class obtainable from an imageset (#201)

* Add Spectrum as a read-only class obtainable from an imageset

Details:
- Add Spectrum C++ object that implements methods get_energies_eV (eV), get_weights (unitless), and get_weighted_wavelenth (angstroms). Spectrum can only be set from constructor.
- Add get_spectrum method to Format and FormatMultiImage
- Add get_spectrum method to ImageSet. Does not cache it like it does for detector, beam, etc.
- Implement get_spectrum in NeXus.  Specifically implements nexusformat/definitions#717, using optional variants to specify a calibrated wavelength and a spectrum.
- Add test for c++ spectrum object
- Add bandwidth calculation. Computes: emin_ev = 1% on CDF and emax_ev = 99% on CDF so 98% or more of the
spectrum is within the range. Added a test for this, which has a computed approximation to a real SwissFEL spectrum (so could be useful in other tests)
- Add weighted energy variance. Recover the weighted variance around the weighted mean. Add test which makes
a clean Gaussian and recovers the same parameters.

Co-authored-by: Graeme Winter <graeme.winter@gmail.com>

SConscript

@@ -153,6 +153,7 @@ if not env_etc.no_boost_python and hasattr(env_etc, "boost_adaptbx_include"):
         target="#/lib/dxtbx_model_ext",
         source=[
             "model/boost_python/beam.cc",
+            "model/boost_python/spectrum.cc",
             "model/boost_python/goniometer.cc",
             "model/boost_python/kappa_goniometer.cc",
             "model/boost_python/multi_axis_goniometer.cc",

format/Format.py

@@ -497,28 +497,34 @@ def _end(self):
 
     def _goniometer(self):
         """Overload this method to read the image file however you like so
-        long as the result is an goniometer."""
+        long as the result is a goniometer."""
         return None
 
     def _detector(self):
         """Overload this method to read the image file however you like so
-        long as the result is an detector."""
+        long as the result is a detector."""
         return None
 
     def _beam(self):
         """Overload this method to read the image file however you like so
-        long as the result is an beam."""
+        long as the result is a beam."""
         return None
 
     def _scan(self):
         """Overload this method to read the image file however you like so
-        long as the result is an scan."""
+        long as the result is a scan."""
         return None
 
     def get_static_mask(self):
         """Overload this method to override the static mask."""
         return None
 
+    def get_spectrum(self):
+        """ Overload this method to read the image file however you like so
+        long as the result is a spectrum
+        """
+        return None
+
     def get_goniometer_shadow_masker(self, goniometer=None):
         """Overload this method to allow generation of dynamic goniometer shadow
         masks to be used during spotfinding or integration."""

format/FormatHDF5EigerNearlyNexus.py

@@ -380,7 +380,7 @@ def _detector(self):
         return self._detector_model
 
     def _beam(self, index=None):
-        self._beam_model = self._beam_factory.load_model(index)
+        self._beam_model, _ = self._beam_factory.read_models(index)
         return self._beam_model
 
     def _scan(self):

format/FormatMultiImage.py

@@ -70,6 +70,9 @@ def get_detector(self, index=None):
     def get_beam(self, index=None):
         return self._beam_instance
 
+    def get_spectrum(self, index=None):
+        raise NotImplementedError
+
     def get_scan(self, index=None):
         return self._scan_instance
 

format/FormatNexus.py

@@ -99,7 +99,7 @@ def _detector(self):
         return self._detector_model
 
     def _beam(self, index=None):
-        self._beam_model = self._beam_factory.load_model(index)
+        self._beam_model, _ = self._beam_factory.read_models(index)
         return self._beam_model
 
     def _scan(self):
@@ -114,6 +114,10 @@ def get_detector(self, index=None):
     def get_beam(self, index=None):
         return self._beam(index)
 
+    def get_spectrum(self, index=None):
+        self._beam_model, _ = self._beam_factory.read_models(index)
+        return self._beam_factory.spectrum
+
     def get_scan(self, index=None):
         if index is None:
             return self._scan()

format/FormatNexusJungfrauHack.py

@@ -223,10 +223,6 @@ def _goniometer(self):
     def _detector(self):
         return self._detector_model
 
-    def _beam(self, index=None):
-        self._beam_model = self._beam_factory.load_model(index)
-        return self._beam_model
-
     def _scan(self):
         return self._scan_model
 

format/nexus.py

@@ -19,6 +19,7 @@
 import dxtbx.model
 from dxtbx.model import (
     Beam,
+    Spectrum,
     Crystal,
     Detector,
     Panel,
@@ -542,36 +543,87 @@ def __init__(self, obj):
         self.obj = obj
         self.model = None
         self.index = None
+        self.spectrum = None
+
+    def read_models(self, index=None):
+        self.load_model(index)
+        return self.model, self.spectrum
 
     def load_model(self, index=None):
         # Cached model
         if self.model is not None and index == self.index:
             return self.model
 
         # Get the items from the NXbeam class
-        wavelength = self.obj.handle["incident_wavelength"]
-        wavelength_weights = self.obj.handle.get("incident_wavelength_weights")
-        if wavelength.shape in (tuple(), (1,)):
-            wavelength_value = wavelength[()]
-        elif len(wavelength.shape) == 1:
-            if wavelength_weights is None:
-                if index is None:
-                    index = 0
-                wavelength_value = wavelength[index]
+        primary_key = "incident_wavelength"
+        wavelength = self.obj.handle[primary_key]
+        spectrum_wavelengths = wavelength
+        spectrum_weights = self.obj.handle.get(primary_key + "_weight")
+
+        # If the wavelength array does not represent spectra, look for spectra
+        # in the variant chain
+        variant_test = wavelength
+        has_variant_spectra = False
+        while spectrum_weights is None:
+            if "variant" in variant_test.attrs:
+                variant_key = variant_test.attrs["variant"]
+                variant_wavelengths = self.obj.handle[variant_key]
+                variant_weights = self.obj.handle.get(variant_key + "_weight")
+                if variant_weights is None:
+                    variant_test = variant_wavelengths  # Keep looking
+                else:
+                    # Found spectra
+                    spectrum_wavelengths = variant_wavelengths
+                    spectrum_weights = variant_weights  # cause while loop to end
+                    has_variant_spectra = True
             else:
-                raise NotImplementedError("Spectra not implemented")
+                break
+
+        if index is None:
+            index = 0
+        self.index = index
+
+        def get_wavelength(wavelength):
+            if wavelength.shape in (tuple(), (1,)):
+                wavelength_value = wavelength[()]
+            else:
+                wavelength_value = wavelength[index]
+            wavelength_units = wavelength.attrs["units"]
+            wavelength_value = float(
+                convert_units(wavelength_value, wavelength_units, "angstrom")
+            )
+            return wavelength_value
+
+        if spectrum_weights is None:
+            # Construct the beam model
+            wavelength_value = get_wavelength(wavelength)
+            self.model = Beam(direction=(0, 0, 1), wavelength=wavelength_value)
         else:
-            raise NotImplementedError("Spectra not implemented")
-        wavelength_units = wavelength.attrs["units"]
+            self.model = Beam()
+            self.model.set_direction((0, 0, 1))
 
-        # Convert wavelength to Angstroms
-        wavelength_value = float(
-            convert_units(wavelength_value, wavelength_units, "angstrom")
-        )
+            wavelength_units = spectrum_wavelengths.attrs["units"]
 
-        # Construct the beam model
-        self.index = index
-        self.model = Beam(direction=(0, 0, 1), wavelength=wavelength_value)
+            if len(spectrum_wavelengths.shape) > 1:
+                spectrum_wavelengths = spectrum_wavelengths[index]
+            else:
+                spectrum_wavelengths = spectrum_wavelengths[()]
+            if len(spectrum_weights.shape) > 1:
+                spectrum_weights = spectrum_weights[index]
+            else:
+                spectrum_weights = spectrum_weights[()]
+
+            spectrum_wavelengths = convert_units(
+                spectrum_wavelengths, wavelength_units, "angstrom"
+            )
+            spectrum_energies = 12398.4187 / spectrum_wavelengths
+            self.spectrum = Spectrum(spectrum_energies, spectrum_wavelengths)
+
+            if has_variant_spectra:
+                wavelength_value = get_wavelength(wavelength)
+                self.model.set_wavelength(wavelength_value)
+            else:
+                self.model.set_wavelength(self.spectrum.get_weighted_wavelength())
         return self.model
 
 

imageset.py

@@ -102,6 +102,19 @@ def get_format_class(self):
         """Get format class name"""
         return self.data().get_format_class()
 
+    def get_spectrum(self, index):
+        """Get the spectrum if available"""
+        kwargs = self.params()
+        if self.data().has_single_file_reader():
+            format_instance = self.get_format_class().get_instance(
+                self.data().get_master_path(), **kwargs
+            )
+        else:
+            format_instance = self.get_format_class().get_instance(
+                self.get_path(index), **kwargs
+            )
+        return format_instance.get_spectrum(self.indices()[index])
+
     def params(self):
         """Get the parameters"""
         return self.data().get_params()

model/__init__.py

@@ -50,6 +50,7 @@
     Scan,
     ScanBase,
     SimplePxMmStrategy,
+    Spectrum,
     VirtualPanel,
     VirtualPanelFrame,
     get_mod2pi_angles_in_range,
@@ -93,6 +94,7 @@
     "ScanBase",
     "ScanFactory",
     "SimplePxMmStrategy",
+    "Spectrum",
     "VirtualPanel",
     "VirtualPanelFrame",
     "get_mod2pi_angles_in_range",

model/boost_python/model_ext.cc

@@ -28,6 +28,7 @@ namespace dxtbx { namespace model { namespace boost_python {
   void export_pixel_to_millimeter();
   void export_experiment();
   void export_experiment_list();
+  void export_spectrum();
 
   BOOST_PYTHON_MODULE(dxtbx_model_ext) {
     export_beam();
@@ -43,6 +44,7 @@ namespace dxtbx { namespace model { namespace boost_python {
     export_pixel_to_millimeter();
     export_experiment();
     export_experiment_list();
+    export_spectrum();
   }
 
 }}}  // namespace dxtbx::model::boost_python

model/boost_python/spectrum.cc

@@ -0,0 +1,83 @@
+/*
+ * spectrum.cc
+ */
+#include <boost/python.hpp>
+#include <boost/python/def.hpp>
+#include <string>
+#include <sstream>
+#include <dxtbx/model/spectrum.h>
+#include <dxtbx/model/boost_python/to_from_dict.h>
+#include <scitbx/array_family/boost_python/flex_wrapper.h>
+
+namespace dxtbx { namespace model { namespace boost_python {
+
+  using namespace boost::python;
+  using scitbx::deg_as_rad;
+  using scitbx::rad_as_deg;
+
+  std::string spectrum_to_string(const Spectrum &spectrum) {
+    std::stringstream ss;
+    ss << spectrum;
+    return ss.str();
+  }
+
+  struct SpectrumPickleSuite : boost::python::pickle_suite {
+    static boost::python::tuple getinitargs(const Spectrum &obj) {
+      return boost::python::make_tuple(obj.get_energies_eV(), obj.get_weights());
+    }
+
+    static boost::python::tuple getstate(boost::python::object obj) {
+      return boost::python::make_tuple(obj.attr("__dict__"));
+    }
+
+    static void setstate(boost::python::object obj, boost::python::tuple state) {
+      DXTBX_ASSERT(boost::python::len(state) == 2);
+
+      // restore the object's __dict__
+      boost::python::dict d =
+        boost::python::extract<boost::python::dict>(obj.attr("__dict__"))();
+      d.update(state[0]);
+    }
+
+    static bool getstate_manages_dict() {
+      return true;
+    }
+  };
+
+  template <>
+  boost::python::dict to_dict<Spectrum>(const Spectrum &obj) {
+    boost::python::dict result;
+    result["energies"] = obj.get_energies_eV();
+    result["weights"] = obj.get_weights();
+    return result;
+  }
+
+  template <>
+  Spectrum *from_dict<Spectrum>(boost::python::dict obj) {
+    Spectrum *s = new Spectrum(boost::python::extract<vecd>(obj["energies"]),
+                               boost::python::extract<vecd>(obj["weights"]));
+    return s;
+  }
+
+  void export_spectrum() {
+    // Export Spectrum
+    class_<Spectrum, boost::shared_ptr<Spectrum> >("Spectrum")
+      .def(init<const Spectrum &>())
+      .def(init<vecd, vecd>((arg("energies"), arg("weights"))))
+      .def("get_energies_eV", &Spectrum::get_energies_eV)
+      .def("get_weights", &Spectrum::get_weights)
+      .def("get_weighted_energy_eV", &Spectrum::get_weighted_energy_eV)
+      .def("get_weighted_energy_variance", &Spectrum::get_weighted_energy_variance)
+      .def("get_weighted_wavelength", &Spectrum::get_weighted_wavelength)
+      .def("get_emin_eV", &Spectrum::get_emin_eV)
+      .def("get_emax_eV", &Spectrum::get_emax_eV)
+      .def("__str__", &spectrum_to_string)
+      .def("to_dict", &to_dict<Spectrum>)
+      .def("from_dict", &from_dict<Spectrum>, return_value_policy<manage_new_object>())
+      .staticmethod("from_dict")
+      .def_pickle(SpectrumPickleSuite());
+
+    scitbx::af::boost_python::flex_wrapper<Spectrum>::plain("flex_Spectrum");
+  }
+
+}}}  // namespace dxtbx::model::boost_python