From ef5f9abab7c08d46f2bff43ad497e44fe3604048 Mon Sep 17 00:00:00 2001 From: David McDonagh <60879630+toastisme@users.noreply.github.com> Date: Wed, 2 Aug 2023 14:15:44 +0100 Subject: [PATCH] Add polychromatic beam (#621) Add new Beam class "PolychromaticBeam" for polychromatic/multi-wavelength/wide bandpass experiments. --------- Co-authored-by: DiamondLightSource-build-server Co-authored-by: Nicholas Devenish --- newsfragments/621.feature | 1 + src/dxtbx/dxtbx_model_ext.pyi | 29 ++++ src/dxtbx/model/__init__.py | 3 + src/dxtbx/model/beam.h | 191 +++++++++++++++++++++++++-- src/dxtbx/model/beam.py | 145 +++++++++++++------- src/dxtbx/model/boost_python/beam.cc | 130 ++++++++++++++++++ tests/model/test_beam.py | 94 ++++++++++++- 7 files changed, 535 insertions(+), 58 deletions(-) create mode 100644 newsfragments/621.feature diff --git a/newsfragments/621.feature b/newsfragments/621.feature new file mode 100644 index 000000000..394ff8e74 --- /dev/null +++ b/newsfragments/621.feature @@ -0,0 +1 @@ +Add new Beam class "PolychromaticBeam" for polychromatic/multi-wavelength/wide bandpass experiments. diff --git a/src/dxtbx/dxtbx_model_ext.pyi b/src/dxtbx/dxtbx_model_ext.pyi index 4cbee00fb..4eb5e1530 100644 --- a/src/dxtbx/dxtbx_model_ext.pyi +++ b/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: ... diff --git a/src/dxtbx/model/__init__.py b/src/dxtbx/model/__init__.py index 26a2fc4d8..a7e3ada72 100644 --- a/src/dxtbx/model/__init__.py +++ b/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", diff --git a/src/dxtbx/model/beam.h b/src/dxtbx/model/beam.h index 9903341ec..a8d72f33c 100644 --- a/src/dxtbx/model/beam.h +++ b/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 get_s0() const { + virtual vec3 get_s0() const { DXTBX_ASSERT(wavelength_ != 0.0); return -direction_ * 1.0 / wavelength_; } - void set_s0(vec3 s0) { + virtual void set_s0(vec3 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,8 +375,7 @@ namespace dxtbx { namespace model { friend std::ostream &operator<<(std::ostream &os, const Beam &b); - private: - double wavelength_; + protected: vec3 direction_; double divergence_; double sigma_divergence_; @@ -384,6 +383,9 @@ namespace dxtbx { namespace model { double polarization_fraction_; double flux_; double transmission_; + + private: + double wavelength_; scitbx::af::shared > s0_at_scan_points_; }; @@ -403,6 +405,171 @@ namespace dxtbx { namespace model { return os; } + class PolychromaticBeam : public Beam { + public: + PolychromaticBeam() { + set_direction(vec3(0.0, 0.0, 1.0)); + set_divergence(0.0); + set_sigma_divergence(0.0); + set_polarization_normal(vec3(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 direction) { + DXTBX_ASSERT(direction.length() > 0); + direction_ = direction.normalize(); + set_divergence(0.0); + set_sigma_divergence(0.0); + set_polarization_normal(vec3(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 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(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 direction, + double divergence, + double sigma_divergence, + vec3 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 get_s0() const { + throw DXTBX_ERROR("PolychromaticBeam has no fixed s0"); + return vec3(0., 0., 0.); + } + + void set_s0(vec3 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 > &s0) { + throw DXTBX_ERROR("PolychromaticBeam has no fixed s0"); + } + + scitbx::af::shared > get_s0_at_scan_points() const { + throw DXTBX_ERROR("PolychromaticBeam has no fixed s0"); + return scitbx::af::shared >(1, (0., 0., 0.)); + } + + vec3 get_s0_at_scan_point(std::size_t index) const { + throw DXTBX_ERROR("PolychromaticBeam has no fixed s0"); + return vec3(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 diff --git a/src/dxtbx/model/beam.py b/src/dxtbx/model/beam.py index 870dd9c4b..7c94cffa7 100644 --- a/src/dxtbx/model/beam.py +++ b/src/dxtbx/model/beam.py @@ -1,15 +1,18 @@ from __future__ import annotations import math +from typing import Tuple import pycbf import libtbx.phil try: - from ..dxtbx_model_ext import Beam + from ..dxtbx_model_ext import Beam, PolychromaticBeam except ModuleNotFoundError: - from dxtbx_model_ext import Beam # type: ignore + from dxtbx_model_ext import Beam, PolychromaticBeam # type: ignore + +Vec3Float = Tuple[float, float, float] beam_phil_scope = libtbx.phil.parse( """ @@ -17,6 +20,10 @@ .expert_level = 1 .short_caption = "Beam overrides" { + type = *monochromatic polychromatic + .type = choice + .help = "Override the beam type" + .short_caption = "beam_type" wavelength = None .type = float @@ -27,6 +34,14 @@ .help = "Override the sample to source direction" .short_caption = "Sample to source direction" + divergence = None + .type = float + .help = "Override the beam divergence" + + sigma_divergence = None + .type = float + .help = "Override the beam sigma divergence" + polarization_normal = None .type = floats(size=3) .help = "Override the polarization normal" @@ -57,62 +72,77 @@ class BeamFactory: will be used, otherwise simplified descriptions can be applied.""" @staticmethod - def from_phil(params, reference=None): + def from_phil( + params: libtbx.phil.scope_extract, + reference: Beam | PolychromaticBeam = None, + ) -> Beam | PolychromaticBeam: """ Convert the phil parameters into a beam model - """ + # Check the input if reference is None: - beam = Beam() + beam = ( + PolychromaticBeam() if params.beam.type == "polychromatic" else Beam() + ) else: beam = reference # Set the parameters - if params.beam.wavelength is not None: - beam.set_wavelength(params.beam.wavelength) - elif reference is None: - raise RuntimeError("No wavelength set") + if params.beam.type == "monochromatic": + if params.beam.wavelength is not None: + beam.set_wavelength(params.beam.wavelength) + elif reference is None: + raise RuntimeError("No wavelength set") if params.beam.direction is not None: beam.set_direction(params.beam.direction) elif reference is None: raise RuntimeError("No beam direction set") + + if params.beam.divergence is not None: + beam.set_divergence(params.beam.divergence) + if params.beam.sigma_divergence is not None: + beam.set_sigma_divergence(params.beam.sigma_divergence) if params.beam.polarization_normal is not None: beam.set_polarization_normal(params.beam.polarization_normal) if params.beam.polarization_fraction is not None: beam.set_polarization_fraction(params.beam.polarization_fraction) + if params.beam.transmission is not None: + beam.set_transmission(params.beam.transmission) + if params.beam.flux is not None: + beam.set_flux(params.beam.flux) return beam @staticmethod - def from_dict(d, t=None): - """Convert the dictionary to a beam model + def from_dict(dict: dict, template: dict = None) -> Beam | PolychromaticBeam: + """Convert the dictionary to a beam model""" - Params: - d The dictionary of parameters - t The template dictionary to use + if template is not None: + if "__id__" in dict and "__id__" in template: + assert ( + dict["__id__"] == template["__id__"] + ), "Beam and template dictionaries are not the same type." - Returns: - The beam model - """ - if d is None and t is None: + if dict is None and template is None: return None - joint = t.copy() if t else {} - joint.update(d) + joint = template.copy() if template else {} + joint.update(dict) # Create the model from the joint dictionary + if joint.get("__id__") == "polychromatic": + return PolychromaticBeam.from_dict(joint) return Beam.from_dict(joint) @staticmethod def make_beam( - sample_to_source=None, - wavelength=None, - s0=None, - unit_s0=None, - divergence=None, - sigma_divergence=None, - ): - + sample_to_source: Vec3Float = None, + wavelength: float = None, + s0: Vec3Float = None, + unit_s0: Vec3Float = None, + divergence: float = None, + sigma_divergence: float = None, + ) -> Beam: if divergence is None or sigma_divergence is None: divergence = 0.0 sigma_divergence = 0.0 @@ -137,19 +167,41 @@ def make_beam( assert s0 return Beam(tuple(map(float, s0))) + @staticmethod + def make_polychromatic_beam( + direction: Vec3Float, + divergence: float = 0.0, + sigma_divergence: float = 0.0, + polarization_normal: Vec3Float = (0.0, 1.0, 0.0), + polarization_fraction: float = 0.5, + flux: float = 0.0, + transmission: float = 1.0, + deg: bool = True, + ) -> PolychromaticBeam: + return PolychromaticBeam( + tuple(map(float, direction)), + float(divergence), + float(sigma_divergence), + tuple(map(float, polarization_normal)), + float(polarization_fraction), + float(flux), + float(transmission), + bool(deg), + ) + @staticmethod def make_polarized_beam( - sample_to_source=None, - wavelength=None, - s0=None, - unit_s0=None, - polarization=None, - polarization_fraction=None, - divergence=None, - sigma_divergence=None, - flux=None, - transmission=None, - ): + sample_to_source: Vec3Float = None, + wavelength: float = None, + s0: Vec3Float = None, + unit_s0: Vec3Float = None, + polarization: Vec3Float = None, + polarization_fraction: float = None, + divergence: float = None, + sigma_divergence: float = None, + flux: float = None, + transmission: float = None, + ) -> Beam: assert polarization assert 0.0 <= polarization_fraction <= 1.0 @@ -199,7 +251,7 @@ def make_polarized_beam( ) @staticmethod - def simple(wavelength): + def simple(wavelength: float) -> Beam: """Construct a beam object on the principle that the beam is aligned with the +z axis, as is quite normal. Also assume the beam has polarization fraction 0.999 and is polarized in the x-z plane, unless @@ -219,7 +271,7 @@ def simple(wavelength): ) @staticmethod - def simple_directional(sample_to_source, wavelength): + def simple_directional(sample_to_source: Vec3Float, wavelength: float) -> Beam: """Construct a beam with direction and wavelength.""" if wavelength > 0.05: @@ -236,8 +288,11 @@ def simple_directional(sample_to_source, wavelength): @staticmethod def complex( - sample_to_source, polarization_fraction, polarization_plane_normal, wavelength - ): + sample_to_source: Vec3Float, + polarization_fraction: float, + polarization_plane_normal: Vec3Float, + wavelength: float, + ) -> Beam: """Full access to the constructor for cases where we do know everything that we need...""" @@ -249,7 +304,7 @@ def complex( ) @staticmethod - def imgCIF(cif_file): + def imgCIF(cif_file: str) -> Beam: """Initialize a detector model from an imgCIF file. N.B. the definition of the polarization plane is not completely helpful in this - it is the angle between the polarization plane and the @@ -263,7 +318,7 @@ def imgCIF(cif_file): return result @staticmethod - def imgCIF_H(cbf_handle): + def imgCIF_H(cbf_handle: pycbf.cbf_handle_struct) -> Beam: """Initialize a detector model from an imgCIF file. N.B. the definition of the polarization plane is not completely helpful in this - it is the angle between the polarization plane and the diff --git a/src/dxtbx/model/boost_python/beam.cc b/src/dxtbx/model/boost_python/beam.cc index bf709d526..45bec9f95 100644 --- a/src/dxtbx/model/boost_python/beam.cc +++ b/src/dxtbx/model/boost_python/beam.cc @@ -25,6 +25,8 @@ namespace dxtbx { namespace model { namespace boost_python { using scitbx::deg_as_rad; using scitbx::rad_as_deg; + // Beam + std::string beam_to_string(const Beam &beam) { std::stringstream ss; ss << beam; @@ -180,6 +182,7 @@ namespace dxtbx { namespace model { namespace boost_python { template <> boost::python::dict to_dict(const Beam &obj) { boost::python::dict result; + result["__id__"] = "monochromatic"; result["direction"] = obj.get_sample_to_source_direction(); result["wavelength"] = obj.get_wavelength(); result["divergence"] = rad_as_deg(obj.get_divergence()); @@ -221,6 +224,102 @@ namespace dxtbx { namespace model { namespace boost_python { return b; } + // PolychromaticBeam + + std::string PolychromaticBeam_to_string(const PolychromaticBeam &beam) { + std::stringstream ss; + ss << beam; + return ss.str(); + } + + struct PolychromaticBeamPickleSuite : boost::python::pickle_suite { + static boost::python::tuple getinitargs(const PolychromaticBeam &obj) { + return boost::python::make_tuple(obj.get_sample_to_source_direction(), + obj.get_divergence(), + obj.get_sigma_divergence(), + obj.get_polarization_normal(), + obj.get_polarization_fraction(), + obj.get_flux(), + obj.get_transmission()); + } + }; + + static PolychromaticBeam *make_PolychromaticBeam(vec3 direction) { + return new PolychromaticBeam(direction); + } + + static PolychromaticBeam *make_PolychromaticBeam_w_divergence(vec3 direction, + double divergence, + double sigma_divergence, + bool deg) { + PolychromaticBeam *beam = NULL; + if (deg) { + beam = new PolychromaticBeam( + direction, deg_as_rad(divergence), deg_as_rad(sigma_divergence)); + } else { + beam = new PolychromaticBeam(direction, divergence, sigma_divergence); + } + return beam; + } + + static PolychromaticBeam *make_PolychromaticBeam_w_all( + vec3 direction, + double divergence, + double sigma_divergence, + vec3 polarization_normal, + double polarization_fraction, + double flux, + double transmission, + bool deg) { + PolychromaticBeam *beam = NULL; + if (deg) { + beam = new PolychromaticBeam(direction, + deg_as_rad(divergence), + deg_as_rad(sigma_divergence), + polarization_normal, + polarization_fraction, + flux, + transmission); + } else { + beam = new PolychromaticBeam(direction, + divergence, + sigma_divergence, + polarization_normal, + polarization_fraction, + flux, + transmission); + } + return beam; + } + + template <> + boost::python::dict to_dict(const PolychromaticBeam &obj) { + boost::python::dict result; + result["__id__"] = "polychromatic"; + result["direction"] = obj.get_sample_to_source_direction(); + result["divergence"] = rad_as_deg(obj.get_divergence()); + result["sigma_divergence"] = rad_as_deg(obj.get_sigma_divergence()); + result["polarization_normal"] = obj.get_polarization_normal(); + result["polarization_fraction"] = obj.get_polarization_fraction(); + result["flux"] = obj.get_flux(); + result["transmission"] = obj.get_transmission(); + return result; + } + + template <> + PolychromaticBeam *from_dict(boost::python::dict obj) { + PolychromaticBeam *b = new PolychromaticBeam( + boost::python::extract >(obj["direction"]), + deg_as_rad(boost::python::extract(obj.get("divergence", 0.0))), + deg_as_rad(boost::python::extract(obj.get("sigma_divergence", 0.0))), + boost::python::extract >( + obj.get("polarization_normal", vec3(0.0, 1.0, 0.0))), + boost::python::extract(obj.get("polarization_fraction", 0.999)), + boost::python::extract(obj.get("flux", 0)), + boost::python::extract(obj.get("transmission", 1))); + return b; + } + void export_beam() { using namespace beam_detail; @@ -305,6 +404,37 @@ namespace dxtbx { namespace model { namespace boost_python { .staticmethod("from_dict") .def_pickle(BeamPickleSuite()); + class_, bases >( + "PolychromaticBeam") + .def("__init__", + make_constructor( + &make_PolychromaticBeam, default_call_policies(), (arg("direction")))) + .def("__init__", + make_constructor(&make_PolychromaticBeam_w_divergence, + default_call_policies(), + (arg("direction"), + arg("divergence"), + arg("sigma_divergence"), + arg("deg") = true))) + .def("__init__", + make_constructor(&make_PolychromaticBeam_w_all, + default_call_policies(), + (arg("direction"), + arg("divergence"), + arg("sigma_divergence"), + arg("polarization_normal"), + arg("polarization_fraction"), + arg("flux"), + arg("transmission"), + arg("deg") = true))) + .def("__str__", &PolychromaticBeam_to_string) + .def("to_dict", &to_dict) + .def("from_dict", + &from_dict, + return_value_policy()) + .staticmethod("from_dict") + .def_pickle(PolychromaticBeamPickleSuite()); + scitbx::af::boost_python::flex_wrapper::plain("flex_Beam"); } diff --git a/tests/model/test_beam.py b/tests/model/test_beam.py index 7f1ef5b9d..2f83a790c 100644 --- a/tests/model/test_beam.py +++ b/tests/model/test_beam.py @@ -5,7 +5,7 @@ from libtbx.phil import parse from scitbx import matrix -from dxtbx.model import Beam +from dxtbx.model import Beam, PolychromaticBeam from dxtbx.model.beam import BeamFactory, beam_phil_scope @@ -176,3 +176,95 @@ def test_beam_object_comparison(): def test_beam_self_serialization(): beam = Beam() assert beam == BeamFactory.from_dict(beam.to_dict()) + + +def test_polychromatic_beam_from_phil(): + params = beam_phil_scope.fetch( + parse( + """ + beam { + type = polychromatic + direction = (0., 0., 1.) + divergence = 0.2 + sigma_divergence = 0.3 + polarization_normal = (0., -1., 0.) + polarization_fraction = .65 + transmission = .5 + flux = .75 + } + """ + ) + ).extract() + + beam = BeamFactory.from_phil(params) + assert isinstance(beam, PolychromaticBeam) + + assert beam.get_sample_to_source_direction() == pytest.approx((0.0, 0.0, 1.0)) + assert beam.get_divergence() == pytest.approx(0.2) + assert beam.get_sigma_divergence() == pytest.approx(0.3) + assert beam.get_polarization_normal() == pytest.approx((0.0, -1.0, 0.0)) + assert beam.get_polarization_fraction() == pytest.approx(0.65) + assert beam.get_transmission() == pytest.approx(0.5) + assert beam.get_flux() == pytest.approx(0.75) + + +def test_polychromatic_beam_from_dict(): + beam = PolychromaticBeam() + assert beam == BeamFactory.from_dict(beam.to_dict()) + + +def test_make_polychromatic_beam(): + + direction = (0.0, 0.0, 1.0) + divergence = 0.2 + sigma_divergence = 0.3 + polarization_normal = (0.0, -1.0, 0.0) + polarization_fraction = 0.65 + transmission = 0.5 + flux = 0.75 + + beam = BeamFactory.make_polychromatic_beam( + direction=direction, + divergence=divergence, + sigma_divergence=sigma_divergence, + polarization_normal=polarization_normal, + polarization_fraction=polarization_fraction, + transmission=transmission, + flux=flux, + ) + + assert beam.get_sample_to_source_direction() == pytest.approx((0.0, 0.0, 1.0)) + assert beam.get_divergence() == pytest.approx(0.2) + assert beam.get_sigma_divergence() == pytest.approx(0.3) + assert beam.get_polarization_normal() == pytest.approx((0.0, -1.0, 0.0)) + assert beam.get_polarization_fraction() == pytest.approx(0.65) + assert beam.get_transmission() == pytest.approx(0.5) + assert beam.get_flux() == pytest.approx(0.75) + + +def test_polychromatic_beam_wavelength_guards(): + beam = PolychromaticBeam() + with pytest.raises(RuntimeError): + _ = beam.get_wavelength() + with pytest.raises(RuntimeError): + _ = beam.get_s0() + with pytest.raises(RuntimeError): + _ = beam.get_num_scan_points() + with pytest.raises(RuntimeError): + _ = beam.get_s0_at_scan_points() + with pytest.raises(RuntimeError): + _ = beam.get_s0_at_scan_point(0) + with pytest.raises(RuntimeError): + beam.reset_scan_points() + with pytest.raises(RuntimeError): + beam.set_wavelength(1.0) + with pytest.raises(RuntimeError): + beam.set_s0((0.0, 0.0, 0.1)) + + +def test_polychromatic_beam_str(): + beam = PolychromaticBeam() + assert ( + beam.__str__() + == "Beam:\n sample to source direction : {0,0,1}\n divergence: 0\n sigma divergence: 0\n polarization normal: {0,1,0}\n polarization fraction: 0.5\n flux: 0\n transmission: 1\n" + )