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