Add probe type to the Beam model (#647)

At the moment this just adds the additional datum and allows
serialisation etc. Old experiment files where this is not set are
interpreted as having Probe.xray. Some electron diffraction format
classes are updated here to record Probe.electron.

Co-authored-by: davidmcdonagh <davidnmcdonagh@gmail.com>

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";