diff --git a/model/__init__.py b/model/__init__.py index 98cd9bcfc..553355024 100644 --- a/model/__init__.py +++ b/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): diff --git a/model/beam.h b/model/beam.h index 65bd92245..93de9cf77 100644 --- a/model/beam.h +++ b/model/beam.h @@ -452,6 +452,145 @@ namespace dxtbx { namespace model { scitbx::af::shared > 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 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 direction, double wavelength) + : Beam(direction, wavelength), + spectrum_energies_(), + spectrum_weights_() {} + + /** + * Initialise all the beam parameters. + * @param direction The beam direction vector. + */ + SpectrumBeam(vec3 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 direction, double wavelength, + double divergence, double sigma_divergence) + : Beam(direction, wavelength, divergence, sigma_divergence), + spectrum_energies_(), + spectrum_weights_() {} + + SpectrumBeam(vec3 direction, double wavelength, + double divergence, double sigma_divergence, + vec3 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 direction, double wavelength, + double divergence, double sigma_divergence, + vec3 polarization_normal, + double polarization_fraction, + double flux, double transmission, + scitbx::af::shared spectrum_energies, scitbx::af::shared 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 spectrum_energies, scitbx::af::shared 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 get_spectrum_energies() const { + return spectrum_energies_; + } + + /** + * Get the spectrum weights + */ + scitbx::af::shared 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(this)->is_similar_to(rhs, + wavelength_tolerance, direction_tolerance, + polarization_normal_tolerance, + polarization_fraction_tolerance)) + return false; + const SpectrumBeam* other = dynamic_cast(&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 spectrum_energies_; + scitbx::af::shared 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 diff --git a/model/beam.py b/model/beam.py index b56845af9..71213baff 100644 --- a/model/beam.py +++ b/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( diff --git a/model/boost_python/beam.cc b/model/boost_python/beam.cc index 02b075424..f2aefe4e5 100644 --- a/model/boost_python/beam.cc +++ b/model/boost_python/beam.cc @@ -29,6 +29,12 @@ namespace dxtbx { namespace model { namespace boost_python { return ss.str(); } + std::string spectrumbeam_to_string(const SpectrumBeam &beam) { + std::stringstream ss; + ss << beam; + return ss.str(); + } + struct BeamPickleSuite : boost::python::pickle_suite { static boost::python::tuple getinitargs(const Beam &obj) { return boost::python::make_tuple(obj.get_sample_to_source_direction(), @@ -67,6 +73,50 @@ namespace dxtbx { namespace model { namespace boost_python { } }; + struct SpectrumBeamPickleSuite : boost::python::pickle_suite { + static + boost::python::tuple getinitargs(const SpectrumBeam &obj) { + return boost::python::make_tuple(obj.get_sample_to_source_direction(), + obj.get_wavelength(), + obj.get_divergence(), + obj.get_sigma_divergence(), + obj.get_polarization_normal(), + obj.get_polarization_fraction(), + obj.get_flux(), + obj.get_transmission(), + obj.get_spectrum_energies(), + obj.get_spectrum_weights()); + } + + static + boost::python::tuple getstate(boost::python::object obj) + { + const SpectrumBeam &beam = boost::python::extract(obj)(); + return boost::python::make_tuple( + obj.attr("__dict__"), + beam.get_s0_at_scan_points()); + } + + static + void setstate(boost::python::object obj, boost::python::tuple state) + { + SpectrumBeam &beam = boost::python::extract(obj)(); + DXTBX_ASSERT(boost::python::len(state) == 2); + + // restore the object's __dict__ + boost::python::dict d = boost::python::extract( + obj.attr("__dict__"))(); + d.update(state[0]); + + // restore the internal state of the C++ object + scitbx::af::const_ref< vec3 > s0_list = boost::python::extract< + scitbx::af::const_ref< vec3 > >(state[1]); + beam.set_s0_at_scan_points(s0_list); + } + + static bool getstate_manages_dict() { return true; } + }; + static Beam *make_beam(vec3 sample_to_source, double wavelength, double divergence, @@ -129,6 +179,67 @@ namespace dxtbx { namespace model { namespace boost_python { return beam; } + static SpectrumBeam *make_spectrumbeam(vec3 sample_to_source, + double wavelength, + double divergence, + double sigma_divergence, + bool deg) { + SpectrumBeam *beam = NULL; + if (deg) { + beam = new SpectrumBeam(sample_to_source, + wavelength, + deg_as_rad(divergence), + deg_as_rad(sigma_divergence)); + } else { + beam = new SpectrumBeam(sample_to_source, wavelength, divergence, sigma_divergence); + } + return beam; + } + + static SpectrumBeam *make_spectrumbeam_w_s0(vec3 s0, + double divergence, + double sigma_divergence, + bool deg) { + SpectrumBeam *beam = NULL; + if (deg) { + beam = new SpectrumBeam(s0, deg_as_rad(divergence), deg_as_rad(sigma_divergence)); + } else { + beam = new SpectrumBeam(s0, divergence, sigma_divergence); + } + return beam; + } + + + static Beam* make_spectrumbeam_w_all(vec3 sample_to_source, + double wavelength, + double divergence, double sigma_divergence, + vec3 polarization_normal, + double polarization_fraction, + double flux, + double transmission, + scitbx::af::shared spectrum_energies, + scitbx::af::shared spectrum_weights, + bool deg) { + SpectrumBeam *beam = NULL; + if (deg) { + beam = new SpectrumBeam(sample_to_source, wavelength, + deg_as_rad(divergence), + deg_as_rad(sigma_divergence), + polarization_normal, + polarization_fraction, + flux, transmission, + spectrum_energies, spectrum_weights); + } else { + beam = new SpectrumBeam(sample_to_source, wavelength, + divergence, sigma_divergence, + polarization_normal, + polarization_fraction, + flux, transmission, + spectrum_energies, spectrum_weights); + } + return beam; + } + static double get_divergence(const Beam &beam, bool deg) { double divergence = beam.get_divergence(); return deg ? rad_as_deg(divergence) : divergence; @@ -301,7 +412,69 @@ namespace dxtbx { namespace model { namespace boost_python { .staticmethod("from_dict") .def_pickle(BeamPickleSuite()); + // Export Beam : BeamBase + class_ , bases > ("SpectrumBeam") + .def(init()) + .def(init()) + .def(init , + double> (( + arg("direction"), + arg("wavelength")))) + .def(init > (( + arg("s0")))) + .def("__init__", + make_constructor( + &make_spectrumbeam, + default_call_policies(), ( + arg("direction"), + arg("wavelength"), + arg("divergence"), + arg("sigma_divergence"), + arg("deg") = true))) + .def("__init__", + make_constructor( + &make_spectrumbeam_w_s0, + default_call_policies(), ( + arg("s0"), + arg("divergence"), + arg("sigma_divergence"), + arg("deg") = true))) + .def("__init__", + make_constructor( + &make_spectrumbeam_w_all, + default_call_policies(), ( + arg("direction"), + arg("wavelength"), + arg("divergence"), + arg("sigma_divergence"), + arg("polarization_normal"), + arg("polarization_fraction"), + arg("flux"), + arg("transmission"), + arg("spectrum_energies"), + arg("spectrum_weights"), + arg("deg") = true))) + .def("get_spectrum_energies", + &SpectrumBeam::get_spectrum_energies) + .def("get_spectrum_weights", + &SpectrumBeam::get_spectrum_weights) + .def("get_weighted_wavelength", + &SpectrumBeam::get_weighted_wavelength) + .def("set_spectrum", + &SpectrumBeam::set_spectrum) + .def("is_similar_to", + &SpectrumBeam::is_similar_to, + (arg("other"), + arg("wavelength_tolerance") = 1e-6, + arg("direction_tolerance") = 1e-6, + arg("polarization_normal_tolerance") = 1e-6, + arg("polarization_fraction_tolerance") = 1e-6, + arg("spectrum_weighted_wavelength_tolerance") = 1e-6)) + .def("__str__", &spectrumbeam_to_string) + .def_pickle(SpectrumBeamPickleSuite()); + scitbx::af::boost_python::flex_wrapper::plain("flex_Beam"); + scitbx::af::boost_python::flex_wrapper::plain("flex_SpectrumBeam"); } }}} // namespace dxtbx::model::boost_python diff --git a/tests/model/test_beam.py b/tests/model/test_beam.py index 9639a24bf..05ebc82af 100644 --- a/tests/model/test_beam.py +++ b/tests/model/test_beam.py @@ -7,8 +7,9 @@ from libtbx.phil import parse from scitbx import matrix -from dxtbx.model import Beam +from dxtbx.model import Beam, SpectrumBeam from dxtbx.model.beam import BeamFactory, beam_phil_scope +from dials.array_family import flex def test_setting_direction_and_wavelength(): @@ -178,3 +179,22 @@ def test_beam_object_comparison(): def test_beam_self_serialization(): beam = Beam() assert beam == BeamFactory.from_dict(beam.to_dict()) + + +def test_spectrum_beam(): + spectrum_energies = flex.double(range(9450, 9550)) + spectrum_weights = flex.double(range(len(spectrum_energies))) + b1 = SpectrumBeam() + b2 = SpectrumBeam() + b1.set_spectrum(spectrum_energies, spectrum_weights) + b2.set_spectrum(spectrum_energies, spectrum_weights) + assert b1.get_weighted_wavelength() == pytest.approx(1.3028567060142213) + assert b1.is_similar_to(b2) + b2.set_spectrum(spectrum_energies + 50, spectrum_weights) + assert not b1.is_similar_to(b2) + + b3 = Beam() + b1.set_wavelength(1.2) + b3.set_wavelength(1.2) + assert b1.is_similar_to(b3) + assert b3.is_similar_to(b1) diff --git a/tests/serialize/test_serialize.py b/tests/serialize/test_serialize.py index 4369f348b..c1bc7ec71 100644 --- a/tests/serialize/test_serialize.py +++ b/tests/serialize/test_serialize.py @@ -7,6 +7,7 @@ from dxtbx.model import ( Beam, + SpectrumBeam, BeamFactory, Goniometer, GoniometerFactory, @@ -36,6 +37,30 @@ def test_beam(): assert b2 != b3 +def test_spectrum_beam(): + b1 = SpectrumBeam((1, 0, 0), 2, 0.1, 0.1) + spectrum_energies = flex.double(range(9450, 9550)) + spectrum_weights = flex.double(range(len(spectrum_energies))) + b1.set_spectrum(spectrum_energies, spectrum_weights) + d = b1.to_dict() + b2 = BeamFactory.from_dict(d) + assert d["direction"] == (1, 0, 0) + assert d["wavelength"] == 2 + assert d["divergence"] == pytest.approx(0.1) + assert d["sigma_divergence"] == pytest.approx(0.1) + assert b1 == b2 + assert "s0_at_scan_points" not in d + + # Test with a template and partial dictionary + d2 = {"direction": (0, 1, 0), "divergence": 0.2} + b3 = BeamFactory.from_dict(d2, d) + assert b3.get_sample_to_source_direction() == (0, 1, 0) + assert b3.get_wavelength() == 2 + assert b3.get_divergence() == pytest.approx(0.2) + assert b3.get_sigma_divergence() == pytest.approx(0.1) + assert b2 != b3 + + def test_beam_with_scan_points(): b1 = Beam((1, 0, 0), 2, 0.1, 0.1)