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/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,29 @@ 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:
+        return std::string("unknown");
+      }
+    }
+
+    void set_probe(Probe probe) {
+      probe_ = probe;
+    }
+
     void reset_scan_points() {
       s0_at_scan_points_.clear();
     }
@@ -324,7 +362,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,
@@ -361,7 +400,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 {
@@ -385,11 +425,13 @@ namespace dxtbx { namespace model {
     double flux_;
     double transmission_;
     scitbx::af::shared<vec3<double> > s0_at_scan_points_;
+    Probe probe_;
   };
 
   /** 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";

src/dxtbx/model/beam.py

@@ -7,9 +7,9 @@
 import libtbx.phil
 
 try:
-    from ..dxtbx_model_ext import Beam
+    from ..dxtbx_model_ext import Beam, Probe
 except ModuleNotFoundError:
-    from dxtbx_model_ext import Beam  # type: ignore
+    from dxtbx_model_ext import Beam, Probe  # type: ignore
 
 beam_phil_scope = libtbx.phil.parse(
     """
@@ -101,6 +101,10 @@ def from_dict(d, t=None):
         joint.update(d)
 
         # Create the model from the joint dictionary
+        if "probe" in joint:
+            joint["probe"] = Probe.values[joint["probe"]]
+        else:
+            joint["probe"] = Probe.xray
         return Beam.from_dict(joint)
 
     @staticmethod
@@ -149,6 +153,7 @@ def make_polarized_beam(
         sigma_divergence=None,
         flux=None,
         transmission=None,
+        probe=Probe.xray,
     ):
         assert polarization
         assert 0.0 <= polarization_fraction <= 1.0
@@ -173,6 +178,7 @@ def make_polarized_beam(
                 float(polarization_fraction),
                 float(flux),
                 float(transmission),
+                probe,
             )
         elif unit_s0:
             assert wavelength
@@ -185,17 +191,23 @@ def make_polarized_beam(
                 float(polarization_fraction),
                 float(flux),
                 float(transmission),
+                probe,
             )
         else:
             assert s0
+            sum_sq_s0 = s0[0] ** 2 + s0[1] ** 2 + s0[2] ** 2
+            assert sum_sq_s0 > 0
+            wavelength = 1.0 / math.sqrt(sum_sq_s0)
             return Beam(
                 tuple(map(float, s0)),
+                wavelength,
                 float(divergence),
                 float(sigma_divergence),
                 tuple(map(float, polarization)),
                 float(polarization_fraction),
                 float(flux),
                 float(transmission),
+                probe,
             )
 
     @staticmethod

src/dxtbx/model/boost_python/beam.cc

@@ -17,6 +17,7 @@
 #include <dxtbx/model/beam.h>
 #include <dxtbx/model/boost_python/to_from_dict.h>
 #include <scitbx/array_family/boost_python/flex_wrapper.h>
+#include <boost/python/enum.hpp>
 
 namespace dxtbx { namespace model { namespace boost_python {
   namespace beam_detail {
@@ -40,7 +41,8 @@ namespace dxtbx { namespace model { namespace boost_python {
                                          obj.get_polarization_normal(),
                                          obj.get_polarization_fraction(),
                                          obj.get_flux(),
-                                         obj.get_transmission());
+                                         obj.get_transmission(),
+                                         obj.get_probe());
       }
 
       static boost::python::tuple getstate(boost::python::object obj) {
@@ -107,6 +109,7 @@ namespace dxtbx { namespace model { namespace boost_python {
                                  double polarization_fraction,
                                  double flux,
                                  double transmission,
+                                 Probe probe,
                                  bool deg) {
       Beam *beam = NULL;
       if (deg) {
@@ -117,7 +120,8 @@ namespace dxtbx { namespace model { namespace boost_python {
                         polarization_normal,
                         polarization_fraction,
                         flux,
-                        transmission);
+                        transmission,
+                        probe);
       } else {
         beam = new Beam(sample_to_source,
                         wavelength,
@@ -126,7 +130,8 @@ namespace dxtbx { namespace model { namespace boost_python {
                         polarization_normal,
                         polarization_fraction,
                         flux,
-                        transmission);
+                        transmission,
+                        probe);
       }
       return beam;
     }
@@ -188,6 +193,7 @@ namespace dxtbx { namespace model { namespace boost_python {
     result["polarization_fraction"] = obj.get_polarization_fraction();
     result["flux"] = obj.get_flux();
     result["transmission"] = obj.get_transmission();
+    result["probe"] = obj.get_probe();
     if (obj.get_num_scan_points() > 0) {
       boost::python::list l;
       scitbx::af::shared<vec3<double> > s0_at_scan_points = obj.get_s0_at_scan_points();
@@ -212,7 +218,8 @@ namespace dxtbx { namespace model { namespace boost_python {
         obj.get("polarization_normal", vec3<double>(0.0, 1.0, 0.0))),
       boost::python::extract<double>(obj.get("polarization_fraction", 0.999)),
       boost::python::extract<double>(obj.get("flux", 0)),
-      boost::python::extract<double>(obj.get("transmission", 1)));
+      boost::python::extract<double>(obj.get("transmission", 1)),
+      boost::python::extract<Probe>(obj.get("probe", Probe::xray)));
     if (obj.has_key("s0_at_scan_points")) {
       boost::python::list s0_at_scan_points =
         boost::python::extract<boost::python::list>(obj["s0_at_scan_points"]);
@@ -224,6 +231,11 @@ namespace dxtbx { namespace model { namespace boost_python {
   void export_beam() {
     using namespace beam_detail;
 
+    enum_<Probe>("Probe")
+      .value("xray", xray)
+      .value("electron", electron)
+      .value("neutron", neutron);
+
     class_<BeamBase, boost::noncopyable>("BeamBase", no_init)
       .def("get_sample_to_source_direction", &BeamBase::get_sample_to_source_direction)
       .def("set_direction", &BeamBase::set_direction)
@@ -254,6 +266,9 @@ namespace dxtbx { namespace model { namespace boost_python {
       .def("set_s0_at_scan_points", &Beam_set_s0_at_scan_points_from_list)
       .def("get_s0_at_scan_points", &BeamBase::get_s0_at_scan_points)
       .def("get_s0_at_scan_point", &BeamBase::get_s0_at_scan_point)
+      .def("get_probe", &BeamBase::get_probe)
+      .def("get_probe_name", &BeamBase::get_probe_name)
+      .def("set_probe", &BeamBase::set_probe)
       .def("reset_scan_points", &BeamBase::reset_scan_points)
       .def("rotate_around_origin",
            &rotate_around_origin,
@@ -298,6 +313,7 @@ namespace dxtbx { namespace model { namespace boost_python {
                              arg("polarization_fraction"),
                              arg("flux"),
                              arg("transmission"),
+                             arg("probe") = Probe::xray,
                              arg("deg") = true)))
       .def("__str__", &beam_to_string)
       .def("to_dict", &to_dict<Beam>)