Aperture Test: Add Drift

This way, we can test if "s" was set to the right value.

examples/aperture/README.rst

@@ -8,10 +8,15 @@ Proton beam undergoing collimation by a rectangular boundary aperture.
 
 We use a 250 MeV proton beam with a horizontal rms beam size of 1.56 mm and a vertical rms beam size of 2.21 mm.
 
-The beam is scraped by a 1 mm x 1.5 mm rectangular aperture.
+After a short drift of 0.123, the beam is scraped by a 1 mm x 1.5 mm rectangular aperture.
 
 In this test, the initial values of :math:`\sigma_x`, :math:`\sigma_y`, :math:`\sigma_t`, :math:`\epsilon_x`, :math:`\epsilon_y`, and :math:`\epsilon_t` must agree with nominal values.
-The test fails if any of the final coordinates for the valid (not lost) particles lie outside the aperture boundary.
+The test fails if:
+
+* any of the final coordinates for the valid (not lost) particles lie outside the aperture boundary or
+* any of the lost particles are inside the aperture boundary or
+* if the sum of lost and kept particles is not equal to the initial particles or
+* if the recorded position :math:`s` for the lost particles does not coincide with the drift distance.
 
 
 Run

examples/aperture/analysis_aperture.py

@@ -53,8 +53,6 @@ def get_moments(beam):
 assert num_particles == len(initial)
 # we lost particles in apertures
 assert num_particles > len(final)
-print(len(final))
-print(len(particles_lost))
 assert num_particles == len(particles_lost) + len(final)
 
 print("Initial Beam:")
@@ -111,3 +109,8 @@ def get_moments(beam):
 print(f"  y_max={particles_lost['position_y'].max()}")
 print(f"  y_min={particles_lost['position_y'].min()}")
 assert np.greater_equal(dy.max(), 0.0)
+
+# check that s is set correctly
+lost_at_s = particles_lost["s_lost"]
+drift_s = np.ones_like(lost_at_s) * 0.123
+assert np.allclose(lost_at_s, drift_s)

examples/aperture/input_aperture.in

@@ -21,12 +21,15 @@ beam.mutpt = 0.0
 ###############################################################################
 # Beamline: lattice elements and segments
 ###############################################################################
-lattice.elements = monitor collimator monitor
+lattice.elements = monitor drift collimator monitor
 lattice.nslice = 1
 
 monitor.type = beam_monitor
 monitor.backend = h5
 
+drift.type = drift
+drift.ds = 0.123
+
 collimator.type = aperture
 collimator.shape = rectangular
 collimator.xmax = 1.0e-3

examples/aperture/run_aperture.py

@@ -50,6 +50,7 @@
 sim.lattice.extend(
     [
         monitor,
+        elements.Drift(0.123),
         elements.Aperture(xmax=1.0e-3, ymax=1.5e-3, shape="rectangular"),
         monitor,
     ]