diff --git a/src/dxtbx/model/beam.py b/src/dxtbx/model/beam.py index 10741b433..5ae08b3d5 100644 --- a/src/dxtbx/model/beam.py +++ b/src/dxtbx/model/beam.py @@ -8,9 +8,9 @@ import libtbx.phil try: - from ..dxtbx_model_ext import Beam, Probe, PolychromaticBeam + from ..dxtbx_model_ext import Beam, PolychromaticBeam, Probe except ModuleNotFoundError: - from dxtbx_model_ext import Beam, Probe, PolychromaticBeam # type: ignore + from dxtbx_model_ext import Beam, PolychromaticBeam, Probe # type: ignore Vec3Float = Tuple[float, float, float] @@ -181,6 +181,7 @@ def make_polychromatic_beam( polarization_fraction: float = 0.5, flux: float = 0.0, transmission: float = 1.0, + probe=Probe.xray, deg: bool = True, ) -> PolychromaticBeam: return PolychromaticBeam( @@ -191,6 +192,7 @@ def make_polychromatic_beam( float(polarization_fraction), float(flux), float(transmission), + probe, bool(deg), ) diff --git a/tests/model/test_beam.py b/tests/model/test_beam.py index 2f83a790c..48a6eb5a1 100644 --- a/tests/model/test_beam.py +++ b/tests/model/test_beam.py @@ -6,7 +6,7 @@ from scitbx import matrix from dxtbx.model import Beam, PolychromaticBeam -from dxtbx.model.beam import BeamFactory, beam_phil_scope +from dxtbx.model.beam import BeamFactory, Probe, beam_phil_scope def test_setting_direction_and_wavelength(): @@ -222,6 +222,7 @@ def test_make_polychromatic_beam(): polarization_fraction = 0.65 transmission = 0.5 flux = 0.75 + probe = Probe.neutron beam = BeamFactory.make_polychromatic_beam( direction=direction, @@ -231,6 +232,7 @@ def test_make_polychromatic_beam(): polarization_fraction=polarization_fraction, transmission=transmission, flux=flux, + probe=probe, ) assert beam.get_sample_to_source_direction() == pytest.approx((0.0, 0.0, 1.0)) @@ -240,6 +242,7 @@ def test_make_polychromatic_beam(): 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) + assert beam.get_probe() == Probe.neutron def test_polychromatic_beam_wavelength_guards(): @@ -266,5 +269,5 @@ 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" + == "Beam:\n probe: x-ray\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" )