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
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 get_from_file(file_path, code_name, preserve_prop_dist=False, **kwargs): x, y, z, px, py, pz, q = dr.read_beam(code_name, file_path, **kwargs) z_avg = np.average(z, weights=q) xi = z - z_avg bunch = ParticleBunch(q, x, y, xi, px, py, pz) if preserve_prop_dist: bunch.prop_distance = z_avg return bunch
def _get_beam_at_specified_time_step_analytically(self, t, beam, g_0, w_0, xi_0, A_x, A_y, phi_x_0, phi_y_0, E, E_p, v_w, K): G = 1 + E / g_0 * t if (G < 1 / g_0).any(): n_part = len(np.where(G < 1 / g_0)[0]) print('Warning: unphysical energy found in {}'.format(n_part) + 'particles due to negative accelerating gradient.') # fix unphysical energies (model does not work well when E<=0) G = np.where(G < 1 / g_0, 1 / g_0, G) phi = 2 * np.sqrt(K * g_0) / E * (G**(1 / 2) - 1) if (E == 0).any(): # apply limit when E->0 idx_0 = np.where(E == 0)[0] phi[idx_0] = np.sqrt(K[idx_0] / g_0[idx_0]) * t[idx_0] A_0 = np.sqrt(A_x**2 + A_y**2) x = A_x * G**(-1 / 4) * np.cos(phi + phi_x_0) v_x = -w_0 * A_x * G**(-3 / 4) * np.sin(phi + phi_x_0) p_x = G * g_0 * v_x / ct.c y = A_y * G**(-1 / 4) * np.cos(phi + phi_y_0) v_y = -w_0 * A_y * G**(-3 / 4) * np.sin(phi + phi_y_0) p_y = G * g_0 * v_y / ct.c delta_xi = (ct.c / (2 * E * g_0) * (G**(-1) - 1) + A_0**2 * K / (2 * ct.c * E) * (G**(-1 / 2) - 1)) xi = xi_0 + delta_xi delta_xi_max = -1 / (2 * E) * (ct.c / g_0 + A_0**2 * K / ct.c) g = (g_0 + E * t + E_p * delta_xi_max * t + E_p / 2 * (ct.c - v_w) * t**2 + ct.c * E_p / (2 * E**2) * np.log(G) + E_p * A_0**2 * K * g_0 / (ct.c * E**2) * (G**(1 / 2) - 1)) p_z = np.sqrt(g**2 - p_x**2 - p_y**2) beam_step = ParticleBunch(beam.q, x, y, xi, p_x, p_y, p_z, prop_distance=beam.prop_distance + t * ct.c) return beam_step
def parray_to_wake_t_beam(p_array): """ Converts an Ocelot ParticleArray to a Wake-T ParticleBunch. Parameters: ----------- p_array : ParticleArray The Ocelot distribution to be converted. Returns: -------- A Wake-T ParticleBunch. """ # Check if Wake-T is installed. if not wake_t_installed: raise ImportError('Wake-T is not installed. ' 'Cannot perform conversion to Wake-T ParticleBunch.') # Get beam matrix and reference values. beam_matrix = p_array.rparticles gamma_ref = p_array.E / m_e_GeV z_ref = p_array.s # Calculate gamma and kinetic momentum (p_kin). dp = beam_matrix[5] b_ref = np.sqrt(1 - gamma_ref**(-2)) gamma = dp*gamma_ref*b_ref + gamma_ref p_kin = np.sqrt(gamma_ref**2 - 1) # Create coordinate arrays in Wake-T units. x = beam_matrix[0] px = beam_matrix[1] * p_kin y = beam_matrix[2] py = beam_matrix[3] * p_kin z = - beam_matrix[4] pz = np.sqrt(gamma**2 - px**2 - py**2 - 1) q = p_array.q_array # Create and return Wake-T distribution. return ParticleBunch(q, x, y, z, px, py, pz, prop_distance=z_ref)
def get_gaussian_bunch_from_twiss(en_x, en_y, a_x, a_y, b_x, b_y, ene, ene_sp, s_t, xi_c, q_tot, n_part, x_off=0, y_off=0, theta_x=0, theta_y=0): """ Creates a 6D Gaussian particle bunch with the specified Twiss parameters. Parameters: ----------- en_x : float Normalized trace-space emittance in the x-plane in units of m*rad. en_y : float Normalized trace-space emittance in the y-plane in units of m*rad. a_x : float Alpha parameter in the x-plane. a_y : float Alpha parameter in the y-plane. b_x : float Beta parameter in the x-plane in units of m. b_y : float Beta parameter in the y-plane in units of m. ene: float Mean bunch energy in non-dimmensional units (beta*gamma). ene_sp: float Relative energy spread in %. s_t: float Bunch duration (standard deviation) in units of fs. xi_c: float Central bunch position in the xi in units of m. q_tot: float Total bunch charge in pC. n_part: int Total number of particles in the bunch. x_off: float Centroid offset in the x-plane in units of m. y_off: float Centroid offset in the y-plane in units of m. theta_x: float Pointing angle in the x-plane in radians. theta_y: float Pointing angle in the y-plane in radians. Returns: -------- A ParticleBunch object. """ # Calculate necessary values n_part = int(n_part) ene_sp = ene_sp / 100 ene_sp_abs = ene_sp * ene s_z = s_t * 1e-15 * ct.c em_x = en_x / ene em_y = en_y / ene g_x = (1 + a_x**2) / b_x g_y = (1 + a_y**2) / b_y s_x = np.sqrt(em_x * b_x) s_y = np.sqrt(em_y * b_y) s_xp = np.sqrt(em_x * g_x) s_yp = np.sqrt(em_y * g_y) p_x = -a_x * em_x / (s_x * s_xp) p_y = -a_y * em_y / (s_y * s_yp) p_x_off = theta_x * ene p_y_off = theta_y * ene q_tot = q_tot / 1e12 # Create normalized gaussian distributions u_x = np.random.standard_normal(n_part) v_x = np.random.standard_normal(n_part) u_y = np.random.standard_normal(n_part) v_y = np.random.standard_normal(n_part) # Calculate transverse particle distributions x = s_x * u_x + x_off xp = s_xp * (p_x * u_x + np.sqrt(1 - np.square(p_x)) * v_x) y = s_y * u_y + y_off yp = s_yp * (p_y * u_y + np.sqrt(1 - np.square(p_y)) * v_y) # Create longitudinal distributions (truncated at -3 and 3 sigma in xi) xi = truncnorm.rvs(-3, 3, loc=xi_c, scale=s_z, size=n_part) pz = np.random.normal(ene, ene_sp_abs, n_part) # Change from slope to momentum and apply offset px = xp * pz + p_x_off py = yp * pz + p_y_off # Charge q = np.ones(n_part) * (q_tot / n_part) return ParticleBunch(q, x, y, xi, px, py, pz)
def track_beam_numerically(self, laser, beam, mode, steps, simulation_code=None, simulation_path=None, time_step=None, auto_update_fields=False, reverse_tracking=False, laser_pos_in_pic_code=None, lon_field=None, lon_field_slope=None, foc_strength=None, field_offset=0, filter_fields=False, filter_sigma=20, laser_evolution=False, laser_z_foc=0, r_max=None, xi_min=None, xi_max=None, n_r=100, n_xi=100, parallel=False, n_proc=None): """ Track the beam through the plasma using a 4th order Runge-Kutta method. Parameters: ----------- laser : LaserPulse Laser used in the plasma stage. beam : ParticleBunch Particle bunch to track. mode : string Mode used to determine the wakefields. Possible values are 'Blowout', 'CustomBlowout', 'FromPICCode' or 'cold_fluid_1d'. steps : int Number of steps in which output should be given. simulation_code : string Name of the simulation code from which fields should be read. Only used if mode='FromPICCode'. simulation_path : string Path to the simulation folder where the fields to read are located. Only used if mode='FromPICCode'. time_step : int Time step at which the fields should be read. auto_update_fields : bool If True, new fields will be read from the simulation folder automatically as the particles travel through the plasma. The first time step will be time_step, and new ones will be loaded as the time of flight of the bunch reaches the next time step available in the simulation folder. reverse_tracking : bool Whether to reverse-track the particles through the stage. Currenly only available for mode= 'FromPICCode'. laser_pos_in_pic_code : float (deprecated) Position of the laser pulse center in the comoving frame in the pic code simulation. lon_field : float Value of the longitudinal electric field at the bunch center at the beginning of the tracking in units of V/m. Only used if mode='CustomBlowout'. lon_field_slope : float Value of the longitudinal electric field slope along z at the bunch center at the beginning of the tracking in units of V/m^2. Only used if mode='CustomBlowout'. foc_strength : float Value of the focusing gradient along the bunch in units of T/m. Only used if mode='CustomBlowout'. field_offset : float If 0, the values of 'lon_field', 'lon_field_slope' and 'foc_strength' will be applied at the bunch center. A value >0 (<0) gives them a positive (negative) offset towards the front (back) of the bunch. Only used if mode='CustomBlowout'. filter_fields : bool If true, a Gaussian filter is applied to smooth the wakefields. This can be useful to remove noise. Only used if mode='FromPICCode'. filter_sigma : float Sigma to be used by the Gaussian filter. laser_evolution : bool If True, the laser pulse transverse profile evolves as a Gaussian in vacuum. If False, the pulse envelope stays fixed throughout the computation. Used only if mode='cold_fluid_1d'. laser_z_foc : float Focal position of the laser along z in meters. It is measured as the distance from the beginning of the PlasmaStage. A negative value implies that the focal point is located before the PlasmaStage. Required only if laser_evolution=True and mode='cold_fluid_1d'. r_max : float Maximum radial position up to which plasma wakefield will be calulated. Required only if mode='cold_fluid_1d'. xi_min : float Minimum longiudinal (speed of light frame) position up to which plasma wakefield will be calulated. Required only if mode='cold_fluid_1d'. xi_max : float Maximum longiudinal (speed of light frame) position up to which plasma wakefield will be calulated. Required only if mode='cold_fluid_1d'. n_r : int Number of grid elements along r to calculate the wakefields. Required only if mode='cold_fluid_1d'. n_xi : int Number of grid elements along xi to calculate the wakefields. Required only if mode='cold_fluid_1d'. parallel : float Determines whether or not to use parallel computation. n_proc : int Number of processes to run in parallel. If None, this will equal the number of physical cores. Returns: -------- A list of size 'steps' containing the beam distribution at each step. """ print('') print('Plasma stage') print('-' * len('Plasma stage')) if mode == "Blowout": WF = BlowoutWakefield(self.n_p, laser) if mode == "CustomBlowout": WF = CustomBlowoutWakefield(self.n_p, laser, np.average(beam.xi, weights=beam.q), lon_field, lon_field_slope, foc_strength, field_offset) elif mode == "FromPICCode": if vpic_installed: WF = WakefieldFromPICSimulation(simulation_code, simulation_path, laser, time_step, self.n_p, filter_fields, filter_sigma, reverse_tracking) else: return [] elif mode == 'cold_fluid_1d': WF = NonLinearColdFluidWakefield(self.calculate_density, laser, laser_evolution, laser_z_foc, r_max, xi_min, xi_max, n_r, n_xi) # Get 6D matrix mat = beam.get_6D_matrix_with_charge() # Plasma length in time t_final = self.length / ct.c t_step = t_final / steps dt = self._get_optimized_dt(beam, WF) iterations = int(t_final / dt) # force at least 1 iteration per step it_per_step = max(int(iterations / steps), 1) iterations = it_per_step * steps dt_adjusted = t_final / iterations # initialize list to store the distribution at each step beam_list = list() # get start time start = time.time() if parallel: # compute in parallel if n_proc is None: num_proc = cpu_count() else: num_proc = n_proc print('Parallel computation in {} processes.'.format(num_proc)) num_part = mat.shape[1] part_per_proc = int(np.ceil(num_part / num_proc)) process_pool = Pool(num_proc) t_s = 0 matrix_list = list() # start computaton try: for p in np.arange(num_proc): matrix_list.append(mat[:, p * part_per_proc:(p + 1) * part_per_proc]) print('') st_0 = "Tracking in {} step(s)... ".format(steps) for s in np.arange(steps): print_progress_bar(st_0, s, steps - 1) if auto_update_fields: WF.check_if_update_fields(s * t_step) partial_solver = partial(runge_kutta_4, WF=WF, dt=dt_adjusted, iterations=it_per_step, t0=s * t_step) matrix_list = process_pool.map(partial_solver, matrix_list) beam_matrix = np.concatenate(matrix_list, axis=1) x, px, y, py, xi, pz, q = beam_matrix new_prop_dist = beam.prop_distance + (s + 1) * t_step * ct.c beam_list.append( ParticleBunch(beam.q, x, y, xi, px, py, pz, prop_distance=new_prop_dist)) finally: process_pool.close() process_pool.join() else: # compute in single process print('Serial computation.') print('') st_0 = "Tracking in {} step(s)... ".format(steps) for s in np.arange(steps): print_progress_bar(st_0, s, steps - 1) beam_matrix = runge_kutta_4(mat, WF=WF, t0=s * t_step, dt=dt_adjusted, iterations=it_per_step) x, px, y, py, xi, pz, q = copy(beam_matrix) new_prop_dist = beam.prop_distance + (s + 1) * t_step * ct.c beam_list.append( ParticleBunch(beam.q, x, y, xi, px, py, pz, prop_distance=new_prop_dist)) # print computing time end = time.time() print("Done ({:1.3f} seconds).".format(end - start)) print('-' * 80) # update beam data last_beam = beam_list[-1] beam.set_phase_space(last_beam.x, last_beam.y, last_beam.xi, last_beam.px, last_beam.py, last_beam.pz) beam.increase_prop_distance(self.length) return beam_list
def track_beam_numerically(self, beam, steps, mode='blowout', laser=None, laser_evolution=False, laser_z_foc=0, r_max=None, xi_min=None, xi_max=None, n_r=100, n_xi=100, parallel=False, n_proc=None): """ Track the beam through the plasma using a 4th order Runge-Kutta method. Parameters: ----------- beam : ParticleBunch Particle bunch to track. steps : int Number of steps in which output should be given. mode: str Mode used to determine the wakefields. Possible values are 'blowout', 'blowout_non_rel' and 'cold_fluid_1d'. laser : LaserPulse Laser used in the plasma stage. Required only if mode='cold_fluid_1d'. laser_evolution : bool If True, the laser pulse transverse profile evolves as a Gaussian in vacuum. If False, the pulse envelope stays fixed throughout the computation. laser_z_foc : float Focal position of the laser along z in meters. It is measured as the distance from the position where the plasma density is n=plasma_dens_top, i.e, as the distance from the beginning (end) of the downramp (upramp). Required only if laser_evolution=True. r_max : float Maximum radial position up to which plasma wakefield will be calulated. Required only if mode='cold_fluid_1d'. xi_min : float Minimum longiudinal (speed of light frame) position up to which plasma wakefield will be calulated. Required only if mode='cold_fluid_1d'. xi_max : float Maximum longiudinal (speed of light frame) position up to which plasma wakefield will be calulated. Required only if mode='cold_fluid_1d'. n_r : int Number of grid elements along r to calculate the wakefields. Required only if mode='cold_fluid_1d'. n_xi : int Number of grid elements along xi to calculate the wakefields. Required only if mode='cold_fluid_1d'. parallel : float Determines whether or not to use parallel computation. n_proc : int Number of processes to run in parallel. If None, this will equal the number of physical cores. Required only if parallel=True. Returns: -------- A list of size 'steps' containing the beam distribution at each step. """ print('') print('Plasma ramp') print('-' * len('Plasma ramp')) if mode == 'blowout': field = PlasmaRampBlowoutField(self.calculate_density) elif mode == 'blowout_non_rel': raise NotImplementedError() elif mode == 'cold_fluid_1d': if self.ramp_type == 'upramp': laser_z_foc = self.length - laser_z_foc field = NonLinearColdFluidWakefield(self.calculate_density, laser, laser_evolution, laser_z_foc, r_max, xi_min, xi_max, n_r, n_xi) # Main beam quantities mat = beam.get_6D_matrix_with_charge() # Plasma length in time t_final = self.length / ct.c t_step = t_final / steps dt = self._get_optimized_dt(beam, field) iterations = int(t_final / dt) # force at least 1 iteration per step it_per_step = max(int(iterations / steps), 1) iterations = it_per_step * steps dt_adjusted = t_final / iterations beam_list = list() start = time.time() if parallel: if n_proc is None: num_proc = cpu_count() else: num_proc = n_proc print('Parallel computation in {} processes.'.format(num_proc)) num_part = mat.shape[1] part_per_proc = int(np.ceil(num_part / num_proc)) process_pool = Pool(num_proc) t_s = 0 matrix_list = list() try: for p in np.arange(num_proc): matrix_list.append(mat[:, p * part_per_proc:(p + 1) * part_per_proc]) print('') st_0 = "Tracking in {} step(s)... ".format(steps) for s in np.arange(steps): print_progress_bar(st_0, s, steps - 1) partial_solver = partial(runge_kutta_4, WF=field, dt=dt_adjusted, iterations=it_per_step, t0=s * t_step) matrix_list = process_pool.map(partial_solver, matrix_list) beam_matrix = np.concatenate(matrix_list, axis=1) x, px, y, py, xi, pz, q = beam_matrix new_prop_dist = beam.prop_distance + (s + 1) * t_step * ct.c beam_list.append( ParticleBunch(beam.q, x, y, xi, px, py, pz, prop_distance=new_prop_dist)) finally: process_pool.close() process_pool.join() else: print('Serial computation.') print('') st_0 = "Tracking in {} step(s)... ".format(steps) for s in np.arange(steps): print_progress_bar(st_0, s, steps - 1) beam_matrix = runge_kutta_4(mat, WF=field, t0=s * t_step, dt=dt_adjusted, iterations=it_per_step) x, px, y, py, xi, pz, q = copy(beam_matrix) new_prop_dist = beam.prop_distance + (s + 1) * t_step * ct.c beam_list.append( ParticleBunch(beam.q, x, y, xi, px, py, pz, prop_distance=new_prop_dist)) end = time.time() print("Done ({:1.3f} seconds).".format(end - start)) print('-' * 80) # update beam data last_beam = beam_list[-1] beam.set_phase_space(last_beam.x, last_beam.y, last_beam.xi, last_beam.px, last_beam.py, last_beam.pz) beam.increase_prop_distance(self.length) return beam_list
def track_beam_numerically(self, beam, steps, non_rel=False, parallel=False, n_proc=None): """Tracks the beam through the plasma lens and returns the final phase space""" print('') print('Plasma lens') print('-' * len('Plasma lens')) if non_rel: field = PlasmaLensField(self.foc_strength) else: field = PlasmaLensFieldRelativistic(self.foc_strength) # Main beam quantities mat = beam.get_6D_matrix_with_charge() # Plasma length in time t_final = self.length / ct.c t_step = t_final / steps dt = self._get_optimized_dt(beam, field) iterations = int(t_final / dt) # force at least 1 iteration per step it_per_step = max(int(iterations / steps), 1) iterations = it_per_step * steps dt_adjusted = t_final / iterations beam_list = list() start = time.time() if parallel: if n_proc is None: num_proc = cpu_count() else: num_proc = n_proc print('Parallel computation in {} processes.'.format(num_proc)) num_part = mat.shape[1] part_per_proc = int(np.ceil(num_part / num_proc)) process_pool = Pool(num_proc) t_s = 0 matrix_list = list() try: for p in np.arange(num_proc): matrix_list.append(mat[:, p * part_per_proc:(p + 1) * part_per_proc]) print('') st_0 = "Tracking in {} step(s)... ".format(steps) for s in np.arange(steps): print_progress_bar(st_0, s, steps - 1) partial_solver = partial(runge_kutta_4, WF=field, dt=dt_adjusted, iterations=it_per_step, t0=s * t_step) matrix_list = process_pool.map(partial_solver, matrix_list) beam_matrix = np.concatenate(matrix_list, axis=1) x, px, y, py, xi, pz, q = beam_matrix new_prop_dist = beam.prop_distance + (s + 1) * t_step * ct.c beam_list.append( ParticleBunch(beam.q, x, y, xi, px, py, pz, prop_distance=new_prop_dist)) finally: process_pool.close() process_pool.join() else: # compute in single process print('Serial computation.') print('') st_0 = "Tracking in {} step(s)... ".format(steps) for s in np.arange(steps): print_progress_bar(st_0, s, steps - 1) beam_matrix = runge_kutta_4(mat, WF=field, t0=s * t_step, dt=dt_adjusted, iterations=it_per_step) x, px, y, py, xi, pz, q = copy(beam_matrix) new_prop_dist = beam.prop_distance + (s + 1) * t_step * ct.c beam_list.append( ParticleBunch(beam.q, x, y, xi, px, py, pz, prop_distance=new_prop_dist)) end = time.time() print("Done ({:1.3f} seconds).".format(end - start)) print('-' * 80) # update beam data last_beam = beam_list[-1] beam.set_phase_space(last_beam.x, last_beam.y, last_beam.xi, last_beam.px, last_beam.py, last_beam.pz) beam.increase_prop_distance(self.length) return beam_list