Add polychromatic beam (#621)

Add new Beam class "PolychromaticBeam" for polychromatic/multi-wavelength/wide bandpass experiments.

---------

Co-authored-by: DiamondLightSource-build-server <DiamondLightSource-build-server@users.noreply.github.com>
Co-authored-by: Nicholas Devenish <ndevenish@gmail.com>

newsfragments/621.feature

@@ -0,0 +1 @@
+Add new Beam class "PolychromaticBeam" for polychromatic/multi-wavelength/wide bandpass experiments.

src/dxtbx/dxtbx_model_ext.pyi

@@ -114,6 +114,35 @@ class Beam(BeamBase):
     def from_dict(data: Dict) -> Beam: ...
     def to_dict(self) -> Dict: ...
 
+class PolychromaticBeam(Beam):
+    @overload
+    def __init__(self, beam: PolychromaticBeam) -> None: ...
+    @overload
+    def __init__(self, direction: Vec3Float) -> None: ...
+    @overload
+    def __init__(
+        self,
+        direction: Vec3Float,
+        divergence: float,
+        sigma_divergence: float,
+        deg: bool = ...,
+    ) -> None: ...
+    @overload
+    def __init__(
+        self,
+        direction: Vec3Float,
+        divergence: float,
+        sigma_divergence: float,
+        polarization_normal: Vec3Float,
+        polarization_fraction: float,
+        flux: float,
+        transmission: float,
+        deg: bool = ...,
+    ) -> None: ...
+    @staticmethod
+    def from_dict(data: Dict) -> PolychromaticBeam: ...
+    def to_dict(self) -> Dict: ...
+
 class CrystalBase:
     @property
     def num_scan_points(self) -> int: ...

src/dxtbx/model/__init__.py

@@ -45,6 +45,7 @@
         OffsetPxMmStrategy,
         Panel,
         ParallaxCorrectedPxMmStrategy,
+        PolychromaticBeam,
         PxMmStrategy,
         Scan,
         ScanBase,
@@ -80,6 +81,7 @@
         OffsetPxMmStrategy,
         Panel,
         ParallaxCorrectedPxMmStrategy,
+        PolychromaticBeam,
         PxMmStrategy,
         Scan,
         ScanBase,
@@ -97,6 +99,7 @@
 __all__ = (
     "Beam",
     "BeamBase",
+    "PolychromaticBeam",
     "BeamFactory",
     "Crystal",
     "CrystalBase",

src/dxtbx/model/beam.h

@@ -188,7 +188,7 @@ namespace dxtbx { namespace model {
       return direction_;
     }
 
-    double get_wavelength() const {
+    virtual double get_wavelength() const {
       return wavelength_;
     }
 
@@ -207,16 +207,16 @@ namespace dxtbx { namespace model {
       direction_ = direction.normalize();
     }
 
-    void set_wavelength(double wavelength) {
+    virtual void set_wavelength(double wavelength) {
       wavelength_ = wavelength;
     }
 
-    vec3<double> get_s0() const {
+    virtual vec3<double> get_s0() const {
       DXTBX_ASSERT(wavelength_ != 0.0);
       return -direction_ * 1.0 / wavelength_;
     }
 
-    void set_s0(vec3<double> s0) {
+    virtual void set_s0(vec3<double> s0) {
       DXTBX_ASSERT(s0.length() > 0);
       direction_ = -s0.normalize();
       wavelength_ = 1.0 / s0.length();
@@ -293,7 +293,7 @@ namespace dxtbx { namespace model {
       s0_at_scan_points_.clear();
     }
 
-    bool operator==(const BeamBase &rhs) const {
+    virtual bool operator==(const BeamBase &rhs) const {
       double eps = 1.0e-6;
 
       // scan-varying model checks
@@ -327,11 +327,11 @@ namespace dxtbx { namespace model {
                   <= eps;
     }
 
-    bool is_similar_to(const BeamBase &rhs,
-                       double wavelength_tolerance,
-                       double direction_tolerance,
-                       double polarization_normal_tolerance,
-                       double polarization_fraction_tolerance) const {
+    virtual bool is_similar_to(const BeamBase &rhs,
+                               double wavelength_tolerance,
+                               double direction_tolerance,
+                               double polarization_normal_tolerance,
+                               double polarization_fraction_tolerance) const {
       // scan varying model checks
       if (get_num_scan_points() != rhs.get_num_scan_points()) {
         return false;
@@ -375,15 +375,17 @@ namespace dxtbx { namespace model {
 
     friend std::ostream &operator<<(std::ostream &os, const Beam &b);
 
-  private:
-    double wavelength_;
+  protected:
     vec3<double> direction_;
     double divergence_;
     double sigma_divergence_;
     vec3<double> polarization_normal_;
     double polarization_fraction_;
     double flux_;
     double transmission_;
+
+  private:
+    double wavelength_;
     scitbx::af::shared<vec3<double> > s0_at_scan_points_;
   };
 
@@ -403,6 +405,171 @@ namespace dxtbx { namespace model {
     return os;
   }
 
+  class PolychromaticBeam : public Beam {
+  public:
+    PolychromaticBeam() {
+      set_direction(vec3<double>(0.0, 0.0, 1.0));
+      set_divergence(0.0);
+      set_sigma_divergence(0.0);
+      set_polarization_normal(vec3<double>(0.0, 1.0, 0.0));
+      set_polarization_fraction(0.5);
+      set_flux(0);
+      set_transmission(1.0);
+    }
+
+    /**
+     * @param direction The beam direction pointing sample to source
+     */
+    PolychromaticBeam(vec3<double> direction) {
+      DXTBX_ASSERT(direction.length() > 0);
+      direction_ = direction.normalize();
+      set_divergence(0.0);
+      set_sigma_divergence(0.0);
+      set_polarization_normal(vec3<double>(0.0, 1.0, 0.0));
+      set_polarization_fraction(0.5);
+      set_flux(0);
+      set_transmission(1.0);
+    }
+
+    /**
+     * @param direction The beam direction pointing sample to source
+     * @param divergence The beam divergence
+     * @param sigma_divergence The standard deviation of the beam divergence
+     */
+    PolychromaticBeam(vec3<double> direction,
+                      double divergence,
+                      double sigma_divergence) {
+      DXTBX_ASSERT(direction.length() > 0);
+      direction_ = direction.normalize();
+      set_divergence(divergence);
+      set_sigma_divergence(sigma_divergence);
+      set_polarization_normal(vec3<double>(0.0, 1.0, 0.0));
+      set_polarization_fraction(0.5);
+      set_flux(0);
+      set_transmission(1.0);
+    }
+
+    /**
+     * @param direction The beam direction pointing sample to source
+     * @param divergence The beam divergence
+     * @param sigma_divergence The standard deviation of the beam divergence
+     * @param polarization_normal The polarization plane
+     * @param polarization_fraction The polarization fraction
+     * @param flux The beam flux
+     * @param transmission The beam transmission
+     */
+    PolychromaticBeam(vec3<double> direction,
+                      double divergence,
+                      double sigma_divergence,
+                      vec3<double> polarization_normal,
+                      double polarization_fraction,
+                      double flux,
+                      double transmission) {
+      DXTBX_ASSERT(direction.length() > 0);
+      direction_ = direction.normalize();
+      set_divergence(divergence);
+      set_sigma_divergence(sigma_divergence);
+      set_polarization_normal(polarization_normal);
+      set_polarization_fraction(polarization_fraction);
+      set_flux(flux);
+      set_transmission(transmission);
+    }
+
+    double get_wavelength() const {
+      throw DXTBX_ERROR("PolychromaticBeam has no fixed wavelength");
+      return -1.;
+    }
+
+    void set_wavelength(double wavelength) {
+      throw DXTBX_ERROR("PolychromaticBeam has no fixed wavelength");
+    }
+
+    vec3<double> get_s0() const {
+      throw DXTBX_ERROR("PolychromaticBeam has no fixed s0");
+      return vec3<double>(0., 0., 0.);
+    }
+
+    void set_s0(vec3<double> s0) {
+      throw DXTBX_ERROR("PolychromaticBeam has no fixed s0");
+    }
+
+    std::size_t get_num_scan_points() const {
+      throw DXTBX_ERROR("PolychromaticBeam has no fixed s0");
+      return 1;
+    }
+
+    void set_s0_at_scan_points(const scitbx::af::const_ref<vec3<double> > &s0) {
+      throw DXTBX_ERROR("PolychromaticBeam has no fixed s0");
+    }
+
+    scitbx::af::shared<vec3<double> > get_s0_at_scan_points() const {
+      throw DXTBX_ERROR("PolychromaticBeam has no fixed s0");
+      return scitbx::af::shared<vec3<double> >(1, (0., 0., 0.));
+    }
+
+    vec3<double> get_s0_at_scan_point(std::size_t index) const {
+      throw DXTBX_ERROR("PolychromaticBeam has no fixed s0");
+      return vec3<double>(0., 0., 0.);
+    }
+
+    void reset_scan_points() {
+      throw DXTBX_ERROR("PolychromaticBeam has no fixed s0");
+    }
+
+    bool operator==(const BeamBase &rhs) const {
+      double eps = 1.0e-6;
+
+      return std::abs(angle_safe(direction_, rhs.get_sample_to_source_direction()))
+               <= eps
+             && std::abs(divergence_ - rhs.get_divergence()) <= eps
+             && std::abs(sigma_divergence_ - rhs.get_sigma_divergence()) <= eps
+             && std::abs(
+                  angle_safe(polarization_normal_, rhs.get_polarization_normal()))
+                  <= eps
+             && std::abs(polarization_fraction_ - rhs.get_polarization_fraction())
+                  <= eps;
+    }
+
+    bool is_similar_to(const BeamBase &rhs,
+                       double wavelength_tolerance,
+                       double direction_tolerance,
+                       double polarization_normal_tolerance,
+                       double polarization_fraction_tolerance) const {
+      return is_similar_to(rhs,
+                           direction_tolerance,
+                           polarization_normal_tolerance,
+                           polarization_fraction_tolerance);
+    }
+
+    bool is_similar_to(const BeamBase &rhs,
+                       double direction_tolerance,
+                       double polarization_normal_tolerance,
+                       double polarization_fraction_tolerance) const {
+      return std::abs(angle_safe(direction_, rhs.get_sample_to_source_direction()))
+               <= direction_tolerance
+             && std::abs(
+                  angle_safe(polarization_normal_, rhs.get_polarization_normal()))
+                  <= polarization_normal_tolerance
+             && std::abs(polarization_fraction_ - rhs.get_polarization_fraction())
+                  <= polarization_fraction_tolerance;
+    }
+  };
+
+  /** Print beam information */
+  inline std::ostream &operator<<(std::ostream &os, const PolychromaticBeam &b) {
+    os << "Beam:\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 << "    flux: " << b.get_flux() << "\n";
+    os << "    transmission: " << b.get_transmission() << "\n";
+    return os;
+  }
+
 }}  // namespace dxtbx::model
 
 #endif  // DXTBX_MODEL_BEAM_H