def track_bunch(self, bunch, steps, backtrack=False):
print('')
print('Drift')
print('-' * len('Drift'))
print("Tracking in {} step(s)... ".format(steps), end='')
l_step = self.length / steps
bunch_list = list()
for i in np.arange(0, steps):
l = (i + 1) * l_step * (1 - 2 * backtrack)
(x, y, xi, px, py, pz) = self._track_step(bunch, l)
new_prop_dist = bunch.prop_distance + l
new_bunch = ParticleBunch(bunch.q,
x,
y,
xi,
px,
py,
pz,
prop_distance=new_prop_dist)
new_bunch.x_ref = bunch.x_ref + l * np.sin(bunch.theta_ref)
new_bunch.theta_ref = bunch.theta_ref
bunch_list.append(new_bunch)
# update bunch data
last_bunch = bunch_list[-1]
bunch.set_phase_space(last_bunch.x, last_bunch.y, last_bunch.xi,
last_bunch.px, last_bunch.py, last_bunch.pz)
bunch.prop_distance = last_bunch.prop_distance
bunch.theta_ref = last_bunch.theta_ref
bunch.x_ref = last_bunch.x_ref
print("Done.")
print('-' * 80)
return bunch_list
def create_new_bunch(self, old_bunch, new_bunch_mat, prop_dist):
q = old_bunch.q
if self.angle != 0:
# angle rotated for prop_dist
theta_step = self.angle * prop_dist / self.length
# magnet bending radius
rho = abs(self.length / self.angle)
# new reference angle and transverse displacement
new_theta_ref = old_bunch.theta_ref + theta_step
sign = -theta_step / abs(theta_step)
new_x_ref = (old_bunch.x_ref + sign * rho *
(np.cos(new_theta_ref) - np.cos(old_bunch.theta_ref)))
else:
# new reference angle and transverse displacement
new_theta_ref = old_bunch.theta_ref
new_x_ref = (old_bunch.x_ref +
self.length * np.sin(old_bunch.theta_ref))
# new prop. distance
new_prop_dist = old_bunch.prop_distance + prop_dist
# rotate distribution if reference angle != 0
if new_theta_ref != 0:
rot = rotation_matrix_xz(new_theta_ref)
new_bunch_mat = np.dot(rot, new_bunch_mat)
new_bunch_mat[0] += new_x_ref
# create new bunch
new_bunch = ParticleBunch(q,
bunch_matrix=new_bunch_mat,
prop_distance=new_prop_dist)
new_bunch.theta_ref = new_theta_ref
new_bunch.x_ref = new_x_ref
return new_bunch