Add SpectrumBeam as derived class from Beam object

model/__init__.py

@@ -26,6 +26,7 @@
 from dxtbx_model_ext import (
     Beam,
     BeamBase,
+    SpectrumBeam,
     Crystal,
     CrystalBase,
     Detector,
@@ -60,6 +61,7 @@
 __all__ = (
     "Beam",
     "BeamBase",
+    "SpectrumBeam",
     "BeamFactory",
     "Crystal",
     "CrystalBase",
@@ -128,6 +130,37 @@ def iter_levelorder(self):
                 queue.extend(node)
 
 
+@boost.python.inject_into(SpectrumBeam)
+class _(object):
+    def to_dict(self):
+        """Convert the SpectrumBeam model to a dictionary
+
+        Returns:
+            A dictionary of the parameters
+        """
+
+        d = super(SpectrumBeam, self).to_dict()
+        d["spectrum_energies"] = tuple(self.get_spectrum_energies())
+        d["spectrum_weights"] = tuple(self.get_spectrum_weights())
+        return d
+
+    def from_dict(d):
+        """Convert the dictionary to a SpectrumBeam model
+
+        Params:
+            d The dictionary of parameters
+
+        Returns:
+            The SpectrumBeam model
+        """
+        beam = SpectrumBeam(Beam.from_dict(d))
+        if "spectrum_energies" in d:
+            beam.set_spectrum(
+                flex.double(d["spectrum_energies"]), flex.double(d["spectrum_weights"])
+            )
+        return beam
+
+
 @boost.python.inject_into(Crystal)
 class _(object):
     def show(self, show_scan_varying=False, out=None):

model/beam.h

@@ -452,6 +452,145 @@ namespace dxtbx { namespace model {
     scitbx::af::shared<vec3<double> > s0_at_scan_points_;
   };
 
+  /** A class to represent a beam with a known spectrum. */
+  class SpectrumBeam : public Beam {
+  public:
+
+    /** Default constructor: initialise all to zero */
+    SpectrumBeam()
+      : Beam(),
+        spectrum_energies_(),
+        spectrum_weights_() {}
+
+    /** Copy constructor */
+    SpectrumBeam(Beam b)
+      : Beam(b),
+        spectrum_energies_(),
+        spectrum_weights_() {}
+
+    /**
+     * Initialise all the beam parameters.
+     * @param direction The beam direction vector.
+     */
+    SpectrumBeam(vec3 <double> s0)
+      : Beam(s0),
+        spectrum_energies_(),
+        spectrum_weights_() {}
+
+    /**
+     * Initialise all the beam parameters. Normalize the direction vector
+     * and give it the length of 1.0 / wavelength
+     * @param wavelength The wavelength of the beam
+     * @param direction The beam direction vector.
+     */
+    SpectrumBeam(vec3 <double> direction, double wavelength)
+      : Beam(direction, wavelength),
+        spectrum_energies_(),
+        spectrum_weights_() {}
+
+    /**
+     * Initialise all the beam parameters.
+     * @param direction The beam direction vector.
+     */
+    SpectrumBeam(vec3 <double> s0, double divergence, double sigma_divergence)
+      : Beam(s0, divergence, sigma_divergence),
+        spectrum_energies_(),
+        spectrum_weights_() {}
+
+    /**
+     * Initialise all the beam parameters. Normalize the direction vector
+     * and give it the length of 1.0 / wavelength
+     * @param wavelength The wavelength of the beam
+     * @param direction The beam direction vector.
+     */
+    SpectrumBeam(vec3 <double> direction, double wavelength,
+         double divergence, double sigma_divergence)
+      : Beam(direction, wavelength, divergence, sigma_divergence),
+        spectrum_energies_(),
+        spectrum_weights_() {}
+
+    SpectrumBeam(vec3 <double> direction, double wavelength,
+         double divergence, double sigma_divergence,
+         vec3<double> polarization_normal,
+         double polarization_fraction,
+         double flux,
+         double transmission)
+      : Beam(direction, wavelength, divergence, sigma_divergence,
+             polarization_normal, polarization_fraction, flux, transmission),
+        spectrum_energies_(),
+        spectrum_weights_() {}
+
+    SpectrumBeam(vec3 <double> direction, 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)
+      : Beam(direction, wavelength, divergence, sigma_divergence,
+             polarization_normal, polarization_fraction, flux, transmission),
+        spectrum_energies_(spectrum_energies), spectrum_weights_(spectrum_weights){}
+
+    /**
+     * 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_;
+    }
+
+    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);
+    }
+
+    /**
+     * Check if two models are similar
+     */
+    bool is_similar_to(const BeamBase &rhs,
+                       double wavelength_tolerance,
+                       double direction_tolerance,
+                       double polarization_normal_tolerance,
+                       double polarization_fraction_tolerance,
+                       double spectrum_weighted_wavelength_tolerance) const {
+      if (!dynamic_cast<const Beam*>(this)->is_similar_to(rhs,
+                                                          wavelength_tolerance, direction_tolerance,
+                                                          polarization_normal_tolerance,
+                                                          polarization_fraction_tolerance))
+        return false;
+      const SpectrumBeam* other = dynamic_cast<const SpectrumBeam* >(&rhs);
+      if (!other) return true;
+      return (std::abs(get_weighted_wavelength() - other->get_weighted_wavelength()) <=
+              spectrum_weighted_wavelength_tolerance);
+    }
+
+  friend std::ostream& operator<<(std::ostream &os, const SpectrumBeam &b);
+
+  private:
+    scitbx::af::shared<double> spectrum_energies_;
+    scitbx::af::shared<double> spectrum_weights_;
+  };
+
   /** Print beam information */
   inline std::ostream &operator<<(std::ostream &os, const Beam &b) {
     os << "Beam:\n";
@@ -466,6 +605,21 @@ namespace dxtbx { namespace model {
     return os;
   }
 
+  /** Print beam information */
+  inline
+  std::ostream& operator<<(std::ostream &os, const SpectrumBeam &b) {
+    os << "SpectrumBeam:\n";
+    os << "    wavelength: " << b.get_wavelength() << "\n";
+    os << "    sample to source direction : "
+       << b.get_sample_to_source_direction().const_ref() << "\n";
+    os << "    divergence: " << b.get_divergence() << "\n";
+    os << "    sigma divergence: " << b.get_sigma_divergence() << "\n";
+    os << "    polarization normal: " << b.get_polarization_normal().const_ref()
+       << "\n";
+    os << "    polarization fraction: " << b.get_polarization_fraction() << "\n";
+    os << "    spectrum weighted wavelength: " << b.get_weighted_wavelength() << "\n";
+    return os;
+  }
 }}  // namespace dxtbx::model
 
 #endif  // DXTBX_MODEL_BEAM_H

model/beam.py

@@ -6,7 +6,7 @@
 import libtbx.phil
 
 import pycbf
-from dxtbx_model_ext import Beam
+from dxtbx_model_ext import Beam, SpectrumBeam
 
 beam_phil_scope = libtbx.phil.parse(
     """
@@ -90,7 +90,10 @@ def from_dict(d, t=None):
         joint.update(d)
 
         # Create the model from the joint dictionary
-        return Beam.from_dict(joint)
+        if "spectrum_energies" in joint:
+            return SpectrumBeam.from_dict(joint)
+        else:
+            return Beam.from_dict(joint)
 
     @staticmethod
     def make_beam(