def test_kirk_patrick_baez():
#beam=Beam.initialize_as_person()
#beam.set_flat_divergence(1e-12, 1e-12)
#beam.x = beam.x*1e-3
#beam.z = beam.z*1e-3
shadow_beam = shadow_source()
beam = Beam()
beam.initialize_from_arrays(shadow_beam.getshonecol(1),
shadow_beam.getshonecol(2),
shadow_beam.getshonecol(3),
shadow_beam.getshonecol(4),
shadow_beam.getshonecol(5),
shadow_beam.getshonecol(6),
shadow_beam.getshonecol(10), 0)
bound1 = BoundaryRectangle(xmax=2.5, xmin=-2.5, ymax=2.5, ymin=-2.5)
bound2 = BoundaryRectangle(xmax=1., xmin=-1., ymax=1., ymin=-1.)
kirk_patrick_baez = CompoundOpticalElement.initialize_as_kirkpatrick_baez(
p=10.,
q=5.,
separation=4.,
theta=89 * np.pi / 180,
bound1=bound1,
bound2=bound2)
beam = kirk_patrick_baez.trace_compound(beam)
beam.plot_good_xz(0)
indices = np.where(beam.flag > 0)
assert_almost_equal(beam.x[indices], 0., 4)
assert_almost_equal(beam.z[indices], 0., 4)
beam.retrace(50.)
beam.plot_good_xz()
print(kirk_patrick_baez.info())
print("Number of good rays: %f" % (beam.number_of_good_rays()))
#beam.histogram()
if do_plot:
plt.show()
