Add probe type to the Beam model

newsfragments/647.feature

@@ -0,0 +1 @@
+The ``Beam`` model now has a ``probe`` value to keep track of the type of radiation.

src/dxtbx/dxtbx_model_ext.pyi

@@ -22,6 +22,8 @@ from scitbx.array_family import shared as flex_shared
 # Attempt to use the stub typing for flex-inheritance
 from scitbx.array_family.flex import FlexPlain
 
+from dxtbx_model_ext import Probe  # type: ignore
+
 # TypeVar for the set of Experiment models that can be joint-accepted
 # - profile, imageset and scalingmodel are handled as 'object'
 TExperimentModel = TypeVar(
@@ -113,6 +115,8 @@ class Beam(BeamBase):
     @staticmethod
     def from_dict(data: Dict) -> Beam: ...
     def to_dict(self) -> Dict: ...
+    @staticmethod
+    def get_probe_from_name(name: str) -> Probe: ...
 
 class PolychromaticBeam(Beam):
     @overload

src/dxtbx/format/FormatGatanDM4.py

@@ -21,6 +21,7 @@
 )
 from dxtbx.format.Format import Format
 from dxtbx.format.FormatMultiImage import FormatMultiImage
+from dxtbx.model.beam import Probe
 
 
 def read_tag(f, byteorder):
@@ -358,6 +359,7 @@ def _beam(self):
             wavelength=wavelength,
             polarization=(0, 1, 0),
             polarization_fraction=0.5,
+            probe=Probe.electron,
         )
 
     def _scan(self):

src/dxtbx/format/FormatMRC.py

@@ -19,6 +19,7 @@
 from dxtbx.format.Format import Format
 from dxtbx.format.FormatMultiImage import FormatMultiImage
 from dxtbx.model import ScanFactory
+from dxtbx.model.beam import Probe
 
 logger = logging.getLogger("dials")
 
@@ -226,6 +227,7 @@ def _beam(self):
             wavelength=wavelength,
             polarization=(0, 1, 0),
             polarization_fraction=0.5,
+            probe=Probe.electron,
         )
 
 

src/dxtbx/format/FormatSMVTimePix_SU.py

@@ -12,6 +12,7 @@
 from scitbx import matrix
 
 from dxtbx.format.FormatSMV import FormatSMV
+from dxtbx.model.beam import Probe
 from dxtbx.model.detector import Detector
 
 
@@ -76,6 +77,7 @@ def _beam(self):
             wavelength=wavelength,
             polarization=(0, 1, 0),
             polarization_fraction=0.5,
+            probe=Probe.electron,
         )
 
     def _scan(self):

src/dxtbx/model/beam.h

@@ -13,6 +13,7 @@
 
 #include <iostream>
 #include <cmath>
+#include <string>
 #include <scitbx/vec3.h>
 #include <scitbx/array_family/shared.h>
 #include <scitbx/array_family/simple_io.h>
@@ -24,6 +25,9 @@ namespace dxtbx { namespace model {
 
   using scitbx::vec3;
 
+  // probe type enumeration
+  enum Probe { xray = 1, electron = 2, neutron = 3 };
+
   /** Base class for beam objects */
   class BeamBase {
   public:
@@ -44,6 +48,8 @@ namespace dxtbx { namespace model {
     virtual std::size_t get_num_scan_points() const = 0;
     virtual scitbx::af::shared<vec3<double> > get_s0_at_scan_points() const = 0;
     virtual vec3<double> get_s0_at_scan_point(std::size_t index) const = 0;
+    virtual Probe get_probe() const = 0;
+    virtual std::string get_probe_name() const = 0;
 
     virtual void set_direction(vec3<double> direction) = 0;
     virtual void set_wavelength(double wavelength) = 0;
@@ -59,6 +65,7 @@ namespace dxtbx { namespace model {
     virtual void set_transmission(double transmission) = 0;
     virtual void set_s0_at_scan_points(
       const scitbx::af::const_ref<vec3<double> > &s0) = 0;
+    virtual void set_probe(Probe probe) = 0;
 
     virtual void reset_scan_points() = 0;
     virtual bool is_similar_to(const BeamBase &rhs,
@@ -82,7 +89,8 @@ namespace dxtbx { namespace model {
           polarization_normal_(0.0, 1.0, 0.0),
           polarization_fraction_(0.999),
           flux_(0),
-          transmission_(1.0) {}
+          transmission_(1.0),
+          probe_(Probe::xray) {}
 
     /**
      * @param s0 The incident beam vector.
@@ -93,7 +101,8 @@ namespace dxtbx { namespace model {
           polarization_normal_(0.0, 1.0, 0.0),
           polarization_fraction_(0.999),
           flux_(0),
-          transmission_(1.0) {
+          transmission_(1.0),
+          probe_(Probe::xray) {
       DXTBX_ASSERT(s0.length() > 0);
       wavelength_ = 1.0 / s0.length();
       direction_ = -s0.normalize();
@@ -110,7 +119,8 @@ namespace dxtbx { namespace model {
           polarization_normal_(0.0, 1.0, 0.0),
           polarization_fraction_(0.999),
           flux_(0),
-          transmission_(1.0) {
+          transmission_(1.0),
+          probe_(Probe::xray) {
       DXTBX_ASSERT(direction.length() > 0);
       direction_ = direction.normalize();
     }
@@ -126,7 +136,8 @@ namespace dxtbx { namespace model {
           polarization_normal_(0.0, 1.0, 0.0),
           polarization_fraction_(0.999),
           flux_(0),
-          transmission_(1.0) {
+          transmission_(1.0),
+          probe_(Probe::xray) {
       DXTBX_ASSERT(s0.length() > 0);
       wavelength_ = 1.0 / s0.length();
       direction_ = -s0.normalize();
@@ -148,7 +159,8 @@ namespace dxtbx { namespace model {
           polarization_normal_(0.0, 1.0, 0.0),
           polarization_fraction_(0.999),
           flux_(0),
-          transmission_(1.0) {
+          transmission_(1.0),
+          probe_(Probe::xray) {
       DXTBX_ASSERT(direction.length() > 0);
       direction_ = direction.normalize();
     }
@@ -162,6 +174,7 @@ namespace dxtbx { namespace model {
      * @param polarization_fraction The polarization fraction
      * @param flux The beam flux
      * @param transmission The beam transmission
+     * @param probe The probe value
      */
     Beam(vec3<double> direction,
          double wavelength,
@@ -170,14 +183,16 @@ namespace dxtbx { namespace model {
          vec3<double> polarization_normal,
          double polarization_fraction,
          double flux,
-         double transmission)
+         double transmission,
+         Probe probe)
         : wavelength_(wavelength),
           divergence_(divergence),
           sigma_divergence_(sigma_divergence),
           polarization_normal_(polarization_normal),
           polarization_fraction_(polarization_fraction),
           flux_(flux),
-          transmission_(transmission) {
+          transmission_(transmission),
+          probe_(probe) {
       DXTBX_ASSERT(direction.length() > 0);
       direction_ = direction.normalize();
     }
@@ -289,6 +304,44 @@ namespace dxtbx { namespace model {
       return s0_at_scan_points_[index];
     }
 
+    Probe get_probe() const {
+      return probe_;
+    }
+
+    std::string get_probe_name() const {
+      // Return a name that matches NeXus definitions from
+      // https://manual.nexusformat.org/classes/base_classes/NXsource.html
+      switch (probe_) {
+      case xray:
+        return std::string("x-ray");
+      case electron:
+        return std::string("electron");
+      case neutron:
+        return std::string("neutron");
+      default:
+        DXTBX_ERROR("Unknown probe type");
+      }
+    }
+
+    static Probe get_probe_from_name(const std::string probe) {
+      // Return a Probe matched to NeXus definitions from
+      // https://manual.nexusformat.org/classes/base_classes/NXsource.html
+
+      if (probe == "x-ray") {
+        return Probe::xray;
+      } else if (probe == "electron") {
+        return Probe::electron;
+      } else if (probe == "neutron") {
+        return Probe::neutron;
+      }
+
+      DXTBX_ERROR("Unknown probe " + probe);
+    }
+
+    void set_probe(Probe probe) {
+      probe_ = probe;
+    }
+
     void reset_scan_points() {
       s0_at_scan_points_.clear();
     }
@@ -324,7 +377,8 @@ namespace dxtbx { namespace model {
                   angle_safe(polarization_normal_, rhs.get_polarization_normal()))
                   <= eps
              && std::abs(polarization_fraction_ - rhs.get_polarization_fraction())
-                  <= eps;
+                  <= eps
+             && (probe_ == rhs.get_probe());
     }
 
     virtual bool is_similar_to(const BeamBase &rhs,
@@ -361,7 +415,8 @@ namespace dxtbx { namespace model {
                   angle_safe(polarization_normal_, rhs.get_polarization_normal()))
                   <= polarization_normal_tolerance
              && std::abs(polarization_fraction_ - rhs.get_polarization_fraction())
-                  <= polarization_fraction_tolerance;
+                  <= polarization_fraction_tolerance
+             && (probe_ == rhs.get_probe());
     }
 
     bool operator!=(const BeamBase &rhs) const {
@@ -383,6 +438,7 @@ namespace dxtbx { namespace model {
     double polarization_fraction_;
     double flux_;
     double transmission_;
+    Probe probe_;
 
   private:
     double wavelength_;
@@ -392,6 +448,7 @@ namespace dxtbx { namespace model {
   /** Print beam information */
   inline std::ostream &operator<<(std::ostream &os, const Beam &b) {
     os << "Beam:\n";
+    os << "    probe: " << b.get_probe_name() << "\n";
     os << "    wavelength: " << b.get_wavelength() << "\n";
     os << "    sample to source direction : "
        << b.get_sample_to_source_direction().const_ref() << "\n";
@@ -415,6 +472,7 @@ namespace dxtbx { namespace model {
       set_polarization_fraction(0.5);
       set_flux(0);
       set_transmission(1.0);
+      set_probe(Probe::xray);
     }
 
     /**
@@ -429,6 +487,7 @@ namespace dxtbx { namespace model {
       set_polarization_fraction(0.5);
       set_flux(0);
       set_transmission(1.0);
+      set_probe(Probe::xray);
     }
 
     /**
@@ -447,6 +506,7 @@ namespace dxtbx { namespace model {
       set_polarization_fraction(0.5);
       set_flux(0);
       set_transmission(1.0);
+      set_probe(Probe::xray);
     }
 
     /**
@@ -457,14 +517,16 @@ namespace dxtbx { namespace model {
      * @param polarization_fraction The polarization fraction
      * @param flux The beam flux
      * @param transmission The beam transmission
+     * @param probe The probe value
      */
     PolychromaticBeam(vec3<double> direction,
                       double divergence,
                       double sigma_divergence,
                       vec3<double> polarization_normal,
                       double polarization_fraction,
                       double flux,
-                      double transmission) {
+                      double transmission,
+                      Probe probe) {
       DXTBX_ASSERT(direction.length() > 0);
       direction_ = direction.normalize();
       set_divergence(divergence);
@@ -473,6 +535,7 @@ namespace dxtbx { namespace model {
       set_polarization_fraction(polarization_fraction);
       set_flux(flux);
       set_transmission(transmission);
+      set_probe(probe);
     }
 
     double get_wavelength() const {
@@ -527,7 +590,8 @@ namespace dxtbx { namespace model {
                   angle_safe(polarization_normal_, rhs.get_polarization_normal()))
                   <= eps
              && std::abs(polarization_fraction_ - rhs.get_polarization_fraction())
-                  <= eps;
+                  <= eps
+             && (probe_ == rhs.get_probe());
     }
 
     bool is_similar_to(const BeamBase &rhs,
@@ -551,13 +615,15 @@ namespace dxtbx { namespace model {
                   angle_safe(polarization_normal_, rhs.get_polarization_normal()))
                   <= polarization_normal_tolerance
              && std::abs(polarization_fraction_ - rhs.get_polarization_fraction())
-                  <= polarization_fraction_tolerance;
+                  <= polarization_fraction_tolerance
+             && (probe_ == rhs.get_probe());
     }
   };
 
   /** Print beam information */
   inline std::ostream &operator<<(std::ostream &os, const PolychromaticBeam &b) {
     os << "Beam:\n";
+    os << "    probe: " << b.get_probe_name() << "\n";
     os << "    sample to source direction : "
        << b.get_sample_to_source_direction().const_ref() << "\n";
     os << "    divergence: " << b.get_divergence() << "\n";