Fixes for tests

cctbx · Aug 2, 2023 · 7f0e289 · 7f0e289
1 parent d2cf031
commit 7f0e289
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git 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
@@ -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"
  )