diff --git a/pyproject.toml b/pyproject.toml index 41263101..777d938e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -66,6 +66,7 @@ markers = [ "StageBasic_linac: Integration tests for StageBasic linacs", "StageQuasistatic2d: Tests for the StageQuasistatic2d class", "StageWakeT: Tests for the StageWakeT class", + "StageHipace: Tets for the StageHipace class", "presets: Tests for collider presets", "impactx: Tests for ImpactX", ] diff --git a/src/abel/classes/stage/impl/stage_hipace.py b/src/abel/classes/stage/impl/stage_hipace.py index 931fe804..0e24b5f4 100644 --- a/src/abel/classes/stage/impl/stage_hipace.py +++ b/src/abel/classes/stage/impl/stage_hipace.py @@ -87,7 +87,7 @@ class StageHipace(Stage): external magnetic field across the beams. If ``True``, the field gradient of the external field is set to enforce the drive beam to undergo an half-interger number of betatron oscillations along the - stage. Defaults to ``False``. + stage. Defaults to ``False``. mesh_refinement : bool, optional Enable HiPACE++ mesh refinement. See the @@ -214,8 +214,8 @@ def __init__(self, length=None, nom_energy_gain=None, plasma_density=None, drive self.no_plasma = no_plasma # external focusing (APL-like) [T/m] + self.driver_half_oscillations = 1.0 self.external_focusing = external_focusing - self._external_focusing_gradient = None # plasma profile self.plasma_profile = SimpleNamespace() @@ -268,13 +268,6 @@ def track(self, beam_incoming, savedepth=0, runnable=None, verbose=False): self._prepare_ramps() self._make_ramp_profile(tmpfolder) - # set external focusing - if self.external_focusing == False: - self._external_focusing_gradient = 0 - if self.external_focusing == True and self._external_focusing_gradient is None: - num_half_oscillations = 1 - self._external_focusing_gradient = self.driver_source.energy/SI.c*(num_half_oscillations*np.pi/self.get_length())**2 # [T/m] - beam0 = beam_incoming driver0 = driver_incoming @@ -350,7 +343,11 @@ def track(self, beam_incoming, savedepth=0, runnable=None, verbose=False): # input file filename_input = 'input_file' path_input = tmpfolder + filename_input - hipace_write_inputs(path_input, filename_beam, filename_driver, self.plasma_density, self.num_steps, time_step, box_range_z, box_size_xy, ion_motion=self.ion_motion, ion_species=self.ion_species, beam_ionization=self.beam_ionization, radiation_reaction=self.radiation_reaction, output_period=output_period, num_cell_xy=self.num_cell_xy, num_cell_z=num_cell_z, driver_only=self.driver_only, density_table_file=density_table_file, no_plasma=self.no_plasma, external_focusing_gradient=self._external_focusing_gradient, mesh_refinement=self.mesh_refinement, do_spin_tracking=self.do_spin_tracking) + + if self.external_focusing and self.external_focusing_gradient is None: + self.external_focusing_gradient = self.calc_external_focusing_gradient() # Set the gradient for external focusing fields if not already set. + + hipace_write_inputs(path_input, filename_beam, filename_driver, self.plasma_density, self.num_steps, time_step, box_range_z, box_size_xy, ion_motion=self.ion_motion, ion_species=self.ion_species, beam_ionization=self.beam_ionization, radiation_reaction=self.radiation_reaction, output_period=output_period, num_cell_xy=self.num_cell_xy, num_cell_z=num_cell_z, driver_only=self.driver_only, density_table_file=density_table_file, no_plasma=self.no_plasma, external_focusing_gradient=self.external_focusing_gradient, mesh_refinement=self.mesh_refinement, do_spin_tracking=self.do_spin_tracking) ## RUN SIMULATION @@ -689,8 +686,819 @@ def get_length(self): #ss = density_table[:,0] return ss.max()-ss.min() return super().get_length() + + + # ================================================== + def matched_beta_function(self, energy_incoming, match_entrance=True, q=SI.e): + ''' + Calculates the matched beta function of the stage. If there is an + upramp, the beta function is magnified by default so that it shrinks to + the correct size when it enters the main flattop plasma stage. Also + takes into account external focusing field B=[g_ext*y, -g_ext*x, 0] if present. + + + Parameters + ---------- + energy_incoming : [eV] float + The energy used for matching. + + match_entrance : bool, optional + Matches the beta function to the upramp or the stage entrance if + ``True``. Otherwise, will match the beta function to the downramp. + Default set to ``True``. + + q : [C] float, optional + Particle charge. Defaults to elementary charge. + + Returns + ------- + beta_function : [m], float + The matched beta function. + ''' + + energy_incoming = energy_incoming*SI.e # [J] + + g = SI.e*self.plasma_density/(2*SI.epsilon_0*SI.c) # [T/m], ion background focusing gradient + if self.external_focusing_gradient is not None: # Add contribution from external field + g = g + self.external_focusing_gradient + + k_beta = np.sqrt(np.abs(q)*g*SI.c/energy_incoming) # [m^-1], betatron wavenumber. + + if match_entrance: + if self.upramp is not None and self.upramp.ramp_beta_mag is not None: + return 1/k_beta * self.upramp.ramp_beta_mag + else: + return 1/k_beta + else: + if self.downramp.ramp_beta_mag is not None: + return 1/k_beta * self.downramp.ramp_beta_mag + else: + raise ValueError('Downramp ramp_beta_mag not defined.') + + + # ============================================= + @property + def external_focusing(self) -> bool: + "Flag for enabling driver guiding using an external magnetic field." + return self._external_focusing + @external_focusing.setter + def external_focusing(self, enable_external_focusing : bool | None): + self._external_focusing = bool(enable_external_focusing) + _external_focusing = False + + + # ============================================= + @property + def external_focusing_gradient(self) -> float: + """ + The focusing gradient g_ext [T/m] for an external azimuthal magnetic + field B = [g_ext*y, -g_ext*x, 0]. + """ + + if self.external_focusing is True and self._external_focusing_gradient is None: + # If external focusing is enabled, but the gradient of the external + # focusing field is not yet set, try to calculate and set it: + + # Check whether the ramps have been set up if they exist + ramps_not_set_up = ( + (self.upramp is not None and self.upramp.length is None) or + (self.downramp is not None and self.downramp.length is None) + ) + + # Check if there is enough information to set up the ramps + can_set_up_ramps = (self.get_nom_energy_gain(ignore_ramps_if_undefined=True) is not None + and self.nom_energy is not None) + + # Make a copy of the stage and set up its ramps if they are not set + # up and can be set up + if ramps_not_set_up and can_set_up_ramps: + stage_copy = copy.deepcopy(self) + stage_copy._prepare_ramps() + else: + stage_copy = self + + if stage_copy.get_length() is not None: + self._external_focusing_gradient = stage_copy.calc_external_focusing_gradient(num_half_oscillations=self.driver_half_oscillations) + + return self._external_focusing_gradient + + @external_focusing_gradient.setter + def external_focusing_gradient(self, g_ext : float | None): + if g_ext is not None and not isinstance(g_ext, float): + raise ValueError("External focusing gradient must be a float or None.") + if self.external_focusing is False: + self._external_focusing_gradient = None + raise ValueError("Cannot set external focusing gradient when self.external_focusing is False.") + self._external_focusing_gradient = g_ext + + _external_focusing_gradient = None + + + # ============================================= + def calc_external_focusing_gradient(self, num_half_oscillations=None, L=None): + """ + Calculate the external focusing gradient g_ext for an azimuthal magnetic + field B = [g_ext*y, -g_ext*x, 0] that gives ``num_half_oscillations`` + half oscillations for the drive beam over the length of the stage. + + Parameters + ---------- + num_half_oscillations : float, optional + Number of half betatron oscillations that the drive beam is + intended to perform. If ``None``, will use ``self.driver_half_oscillations``. + Defaults to ``None``. + + L : [m] float, optional + The length over which the driver will be guided. If ``None``, will + extract the value using ``self.get_length()``. If the stage does + have ramps that have not been fully set up, a deepcopy of the stage + is created to set up its ramps using + :func:`Stage._prepare_ramps() ` so that + the stage total length is defined. + + Returns + ------- + g_ext : [T/m] float + The gradient for the azimuthal magnetic field. + """ + if L is None: + + # Make a copy of the stage and set up its ramps if they are not set up + ramps_not_set_up = ( + (self.upramp is not None and self.upramp.length is None) or + (self.downramp is not None and self.downramp.length is None) + ) + if ramps_not_set_up: + stage_copy = copy.deepcopy(self) + stage_copy._prepare_ramps() + L = stage_copy.get_length() + else: + stage_copy = self + L = stage_copy.get_length() + if L is None: + L = stage_copy.length_flattop # If there are no ramps, can use either length or length_flattop. + + if L is None: + raise ValueError('Stage length is not set.') + + if num_half_oscillations is None: + num_half_oscillations = self.driver_half_oscillations + + return self.driver_source.energy/SI.c*(num_half_oscillations*np.pi/L)**2 # [T/m] + + + # ================================================== + def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, plasma_density=None, driver_half_oscillations=None, q=SI.e): + """ + Calculate the stage length that gives ``num_beta_osc`` betatron + oscillations for a particle with given initial energy ``initial_energy`` + in a uniform plasma stage (excluding ramps) with defined nominal + acceleration gradient and plasma density ``plasma_density``. + + Will take into account the contribution from an external linear magnetic + field B=[g_ext*y, -g_ext*x, 0] if :attr:`self.external_focusing ` + is set to ``True``. + + Parameters + ---------- + num_beta_osc : float + Total number of design betatron oscillations that the electron + should perform through the plasma stage excluding ramps. + + initial_energy : [eV] float, optional + The initial energy of the particle at the start of the plasma stage. + Defaults to ``self.nom_energy``. + + plasma_density : [m^-3] float, optional + The plasma density of the plasma stage. Defaults to + ``self.plasma_density``. + + driver_half_oscillations : float, optional + Number of half betatron oscillations that the drive beam is + intended to perform. If ``None``, will use :attr:`StageHipace.driver_half_oscillations `. + Defaults to ``None``. + + q : [C] float, optional + Particle charge. q * nom_accel_gradient must be positive. Defaults + to elementary charge. + + + Returns + ------- + length : [m] float + Length of the plasma stage excluding ramps matched to the given + number of betatron oscillations. + """ + + from scipy.optimize import fsolve + + if initial_energy is None: + if self.nom_energy is None: + raise ValueError('Stage.nom_energy not set.') + initial_energy = self.nom_energy + initial_energy = initial_energy*SI.e # [J] + + if self.nom_accel_gradient_flattop is None: + raise ValueError('Stage.nom_accel_gradient_flattop not set.') + nom_accel_gradient = self.nom_accel_gradient_flattop + + if q * nom_accel_gradient < 0: + raise ValueError('q * nom_accel_gradient must be positive.') + + if plasma_density is None: + if self.plasma_density is None: + raise ValueError('Stage.plasma_density not set.') + plasma_density = self.plasma_density + + if num_beta_osc < 0: + raise ValueError('Number of input betatron oscillations must be positive.') + + if driver_half_oscillations is None: + driver_half_oscillations = self.driver_half_oscillations + if driver_half_oscillations < 0: + raise ValueError('Number of driver oscillations must be positive.') + + # Assess whether the ramps have been set up + ramps_not_set_up = ( + (self.upramp is not None and self.upramp.length is None) or + (self.downramp is not None and self.downramp.length is None) + ) + + # Make a copy of the stage + stage_copy = copy.deepcopy(self) + + # The function to be used for solving the equation for phase advance numerically + def rhs(L): + g = SI.e*plasma_density/(2*SI.epsilon_0*SI.c) # [T/m], ion background focusing gradient + + # Set up the ramps using the stage copy + if ramps_not_set_up: + stage_copy.length_flattop = L[0] # L is an ndarray with one element + stage_copy._prepare_ramps() + L_ramps = stage_copy.get_ramp_length() + + if self.external_focusing: # Add contribution from external field used for driver guiding + g_ext = stage_copy.calc_external_focusing_gradient(num_half_oscillations=driver_half_oscillations, L=L+L_ramps) + g = g + g_ext + + prefactor = 2*np.sqrt(np.abs(q)*g*SI.c) / (q*nom_accel_gradient) + energy_scaling = np.sqrt(initial_energy + q*nom_accel_gradient*L) - np.sqrt(initial_energy) + return prefactor * energy_scaling + + # Solve 2*np.pi*num_beta_osc = rhs(L) + solution = fsolve(lambda L: rhs(L) - 2*np.pi*num_beta_osc , x0=1) + length = solution[0] + + return length + + + # ================================================== + def match_length_2_num_beta_osc(self, num_beta_osc, driver_half_oscillations=None, set_consistent_params=True, q=SI.e): + """ + Set :attr:`self.length_flattop ` for a + uniform plasma stage such that a particle with initial energy + :attr:`self.nom_energy ` will perform + ``num_beta_osc`` betatron oscillations through the stage (including any + existing ramps). + + - Assumes that each of the (uniform) ramps are configured to give pi/2 + phase advance for the main beam. + + - The stage length calculation is performed using + :meth:`StageHipace.calc_flattop_num_beta_osc() `. + + - Also set :attr:`self._external_focusing_gradient ` + and :attr:`self.driver_half_oscillations ` + to be consistent with the calculated stage length + if ``set_consistent_params`` and set :attr:`self.external_focusing ` + are ``True``. + + + Parameters + ---------- + num_beta_osc : float + Total number of design betatron oscillations that the electron + should perform through the plasma stage excluding ramps. + + driver_half_oscillations : float, optional + Number of half betatron oscillations that the drive beam is + intended to perform. If ``None``, will use :attr:`StageHipace.driver_half_oscillations ` + Defaults to ``None``. + + set_consistent_params : bool, optional + Flag for setting :attr:`self._external_focusing_gradient ` + and :attr:`self.driver_half_oscillations ` + to be consistent with the calculated stage length. Defaults to + ``True``. + q : [C] float, optional + Particle charge. q * nom_accel_gradient must be positive. Defaults + to elementary charge. + + + Returns + ------- + None + """ + + # Assess whether length flattop can be set + if self._length_flattop_calc is not None and self._length_flattop is None: + from abel.classes.stage.stage import VariablesOverspecifiedError + raise VariablesOverspecifiedError("Stage length already known/calculateable, cannot set.") + + if self.has_ramp(): + # Calculate the number of betatron oscillations that the main beam + # should perform in the flattop: + if self.upramp.ramp_shape != 'uniform' or self.downramp.ramp_shape != 'uniform': + raise ValueError('This method assumes uniform ramps.') + if self.upramp.length_flattop is not None or self.downramp.length_flattop is not None: + raise ValueError('This method assumes uniform ramps with length set to give pi/2 phase advance for the main beam. The lengths for the ramps are already set. Setting a new length for the flattop will give wrong results.') + num_beta_osc_flattop = num_beta_osc - 0.5 # The ramps are by default set up to give pi/2 phase advance for the main beam. + else: + num_beta_osc_flattop = num_beta_osc + + if driver_half_oscillations is None: + driver_half_oscillations = self.driver_half_oscillations + + # Calculate the length of the flattop stage + length_flattop = self.calc_length_num_beta_osc(num_beta_osc=num_beta_osc_flattop, + initial_energy=self.nom_energy, + plasma_density=self.plasma_density, + driver_half_oscillations=driver_half_oscillations, + q=q) + + # Set the length of the flattop stage + self.length_flattop = length_flattop + + # Set the external focusing gradient for the driver guiding field and + # the number of half betatron oscillations that the drive beam is + # intended to perform. + if self.external_focusing and set_consistent_params: + # Assess whether the ramps have been set up + ramps_not_set_up = ( + (self.upramp is not None and self.upramp.length is None) or + (self.downramp is not None and self.downramp.length is None) + ) + + if ramps_not_set_up: + # Make a copy of the stage + stage_copy = copy.deepcopy(self) + stage_copy._prepare_ramps() + else: + stage_copy = self + + length = length_flattop + stage_copy.get_ramp_length() + self.external_focusing_gradient = self.calc_external_focusing_gradient(num_half_oscillations=driver_half_oscillations, L=length) + self.driver_half_oscillations = driver_half_oscillations + + + # ============================================= + def driver_guiding_trajectory(self, num_steps=None, dacc_gradient=0.0): + """ + Estimate the trajectory that the drive beam will follow when driver + guiding with an external linear azimuthal magnetic field is applied to a + drive beam with an initial angular offset. The calculations are done by + integrating simplified equations of motion. + + Parameters + ---------- + num_steps : int, optional + Number of time steps. If ``None``, will calculate the number of time + steps such that the step size is a small fraction of the matched + beta function of the drive beam. Defaults to ``None``. + + dacc_gradient : [V/m] float, optional + The decceleration gradient. Drive beam charge * decceleration + gradient must be negative. Defaults to 0.0. + + + Returns + ------- + s_trajectory : [m] 1D float ndarray + Longitudinal coordinate of the drive beam trajectory. Reference is + set at the start of the plasma stage. + + x_trajectory : [m] 1D float ndarray + x-coordinate of the drive beam trajectory. + + y_trajectory : [m] 1D float ndarray + y-coordinate of the drive beam trajectory. + """ + + from abel.utilities.relativity import energy2momentum + from abel.utilities.statistics import weighted_mean + + driver = self.driver_source.track() + + energy_thres = 10*driver.particle_mass*SI.c**2/SI.e # [eV], 10 * particle rest energy. Gives beta=0.995. + pz_thres = energy2momentum(energy_thres, unit='eV', m=driver.particle_mass) + pz0 = energy2momentum(driver.energy(), unit='eV', m=driver.particle_mass) + + if pz0 < pz_thres: + raise ValueError('This estimate is only valid for a relativistic beam.') + + q = driver.particle_charge() # [C], particle charge including charge sign. + if q * dacc_gradient > 0.0: + raise ValueError('Drive beam charge * decceleration gradient must be negative.') + + # Make a copy of the stage and set up its ramps if they are not set up + ramps_not_set_up = ( + (self.upramp is not None and self.upramp.length is None) or + (self.downramp is not None and self.downramp.length is None) + ) + if ramps_not_set_up: + stage_copy = copy.deepcopy(self) + stage_copy._prepare_ramps() + else: + stage_copy = self + + L = stage_copy.get_length() # [m] + + if pz0 + q * dacc_gradient * L/SI.c < pz_thres: + raise ValueError('The energy depletion will be too severe. This estimate is only valid for a relativistic beam.') + + # Get the focusing field gradient + g = self.external_focusing_gradient # [T/m] + if g is None: + g = 0.0 + + # Determine the step size + if num_steps is None: + matched_beta = self.matched_beta_function(driver.energy()) + num_steps = int(L /(matched_beta/20)) + + ds = L/(num_steps-1) # [m], step size + + # Initialise arrays + s_trajectory = np.full(num_steps, None, dtype=float) + x_trajectory = np.full(num_steps, None, dtype=float) + y_trajectory = np.full(num_steps, None, dtype=float) + + # Set initial parameters + prop_length = 0 + x0 = driver.x_offset() + x = x0 + y0 = driver.y_offset() + y = y0 + s_trajectory[0] = prop_length + x_trajectory[0] = x0 + y_trajectory[0] = y0 + px = weighted_mean(driver.pxs(), driver.weightings(), clean=False) + py = weighted_mean(driver.pys(), driver.weightings(), clean=False) + pz = pz0 # Can add option for deceleration using a gradient + + # Solve the equation of motion + i = 0 + while i < num_steps - 1: + + # Drift + prop_length = prop_length + 1/2*ds + x = x + px/pz*1/2*ds + y = y + py/pz*1/2*ds + + # Kick + dpx = q*g*x*ds + px = px + dpx + dpy = q*g*y*ds + py = py + dpy + pz = pz0 + q * dacc_gradient * prop_length/SI.c # dacc_gradient > 0 + + # Drift + prop_length = prop_length + 1/2*ds + x = x + px/pz*1/2*ds + y = y + py/pz*1/2*ds + + i = i + 1 + s_trajectory[i] = prop_length + x_trajectory[i] = x + y_trajectory[i] = y + + return s_trajectory, x_trajectory, y_trajectory + + + # ============================================= + def estimate_beam_trajectory(self, beam, num_steps=None): + """ + Estimate the trajectory for the main beam following the trajectory of a + drive beam generated by ``self.driver_source``. + + Effects such as driver guiding with an external linear azimuthal + magnetic field, background ion focusing and uniform plasma density ramps + are taken into account. The calculations are done by integrating + simplified equations of motion. + + Parameters + ---------- + beam : ``Beam`` + The main beam to be tracked. + + num_steps : int, optional + Number of time steps. If ``None``, will calculate the number of time + steps such that the step size is a small fraction of the matched + beta function of the main beam. Defaults to ``None``. + + + Returns + ------- + s_trajectory : [m] 1D float ndarray + Longitudinal coordinate of the main beam trajectory. Reference is + set at the start of the plasma stage. + + x_trajectory : [m] 1D float ndarray + x-coordinate of the main beam trajectory. + + y_trajectory : [m] 1D float ndarray + y-coordinate of the main beam trajectory. + + px_trajectory : [kg m/s] 1D float ndarray + Mean x-component of the beam momentum along the trajectory. + + py_trajectory : [kg m/s] 1D float ndarray + Mean y-component of the beam momentum along the trajectory. + + pz_trajectory : [kg m/s] 1D float ndarray + Mean longitudinal of the beam momentum along the trajectory. + + driver_x_trajectory : [m] 1D float ndarray, optional + The x-coordinate of trajectory of the drive beam. + + driver_y_trajectory : [m] 1D float ndarray, optional + The y-coordinate of trajectory of the drive beam. + """ + + from abel.utilities.statistics import weighted_mean + from abel.utilities.relativity import energy2momentum + + # Prepare parameters + x0 = beam.x_offset() + y0 = beam.y_offset() + weights = beam.weightings() + px0 = weighted_mean(beam.pxs(), weights, clean=False) + py0 = weighted_mean(beam.pys(), weights, clean=False) + pz0 = energy2momentum(beam.energy(), unit='eV', m=beam.particle_mass) + q = beam.particle_charge() + + L = self.get_length() # [m], total length including any ramps. + if num_steps is None: + matched_beta = self.matched_beta_function(beam.energy()) + num_steps = int(L /(matched_beta/20)) + + # Actual calculations + s_trajectory, x_trajectory, y_trajectory, px_trajectory, py_trajectory, pz_trajectory, driver_x_trajectory, driver_y_trajectory = self._estimate_beam_trajectory(s0=0.0, + x0=x0, + y0=y0, + px0=px0, + py0=py0, + pz0=pz0, + q=q, + num_steps=num_steps, + driver_x_trajectory=None, + driver_y_trajectory=None) + + return s_trajectory, x_trajectory, y_trajectory, px_trajectory, py_trajectory, pz_trajectory, driver_x_trajectory, driver_y_trajectory + + + # ============================================= + def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps, driver_x_trajectory=None, driver_y_trajectory=None): + """ + Helper function for estimating the trajectory for the main beam + following the trajectory of a drive beam defined by + ``driver_x_trajectory`` and ``driver_y_trajectory``. + + Effects such as driver guiding with an external linear azimuthal + magnetic field, background ion focusing and uniform plasma density ramps + are taken into account. The calculations are done by integrating + simplified equations of motion. + + + Parameters + ---------- + s0 : [m] float + The initial longitudinal coordinate of the main beam. + + x0 : [m] float + The intial x-coordinate of the main beam. + + y0 : [m] float + The intial y-coordinate of the main beam. + + px0 : [kg m/s] float + The intial x-momentum of the main beam. + + py0 : [kg m/s] float + The intial y-momentum of the main beam. + + pz0 : [kg m/s] float + The intial z-momentum of the main beam. + + q : [C] flloat + The particle charge of the main beam. + + num_steps : int + Number of time steps. + + driver_x_trajectory : [m] 1D float ndarray, optional + The x-coordinate of trajectory of the drive beam. The length of + ``driver_x_trajectory`` must be the same as ``num_steps``. Is + automatically calculated if ``None``. Defaults to ``None``. + + driver_y_trajectory : [m] 1D float ndarray, optional + The y-coordinate of trajectory of the drive beam. The length of + ``driver_y_trajectory`` must be the same as ``num_steps``. Is + automatically calculated if ``None``. Defaults to ``None``. + + + Returns + ------- + s_trajectory : [m] 1D float ndarray + Longitudinal coordinate of the main beam trajectory. Reference is + set at the start of the plasma stage. + + x_trajectory : [m] 1D float ndarray + x-coordinate of the main beam trajectory. + + y_trajectory : [m] 1D float ndarray + y-coordinate of the main beam trajectory. + + px_trajectory : [kg m/s] 1D float ndarray + Mean x-component of the beam momentum along the trajectory. + + py_trajectory : [kg m/s] 1D float ndarray + Mean y-component of the beam momentum along the trajectory. + + pz_trajectory : [kg m/s] 1D float ndarray + Mean longitudinal of the beam momentum along the trajectory. + + driver_x_trajectory : [m] 1D float ndarray, optional + The x-coordinate of trajectory of the drive beam. + + driver_y_trajectory : [m] 1D float ndarray, optional + The y-coordinate of trajectory of the drive beam. + """ + from abel.utilities.other import find_closest_value_in_arr + + if q * self.nom_accel_gradient_flattop > 0.0: + raise ValueError('Beam charge * self.nom_accel_gradient_flattop gradient must be negative.') + + # Make a copy of the stage and set up its ramps if they are not set up + ramps_not_set_up = ( + (self.upramp is not None and self.upramp.length is None) or + (self.downramp is not None and self.downramp.length is None) + ) + if ramps_not_set_up: + stage_copy = copy.deepcopy(self) + stage_copy._prepare_ramps() + else: + stage_copy = self + + # Calculate the focusing field gradient + g0 = SI.e*stage_copy.plasma_density/(2*SI.epsilon_0*SI.c) # [T/m] + g = g0 + if stage_copy.external_focusing_gradient is not None: + g = g0 + stage_copy.external_focusing_gradient + + # Only calculate the time step size when calling estimate_beam_trajectory() from a flattop + if not stage_copy.is_upramp() and not stage_copy.is_downramp(): + + # Set the step size + L = stage_copy.get_length() # [m], total length including any ramps. + stage_copy.ds = L / (num_steps-1) # [m], step size + + if stage_copy.has_ramp(): + # Set the number of time steps in the upramp and downramp + ss_helper = np.arange(num_steps) * stage_copy.ds + idx_upramp_end, _ = find_closest_value_in_arr(ss_helper, stage_copy.upramp.length) + num_steps_upramp = idx_upramp_end + 1 # Number of time steps for the upramp + idx_flat_end, _ = find_closest_value_in_arr(ss_helper, stage_copy.upramp.length+stage_copy.length_flattop) # Index marking the end of the flattop. + num_steps_downramp = num_steps - idx_flat_end - 1 # Number of time steps for the downramp + stage_copy.idx_flat_end = idx_flat_end + stage_copy.num_steps_upramp = num_steps_upramp + stage_copy.num_steps_downramp = num_steps_downramp + + ds = stage_copy.ds + + # Calculate the drive beam trajectory + if driver_x_trajectory is None or driver_y_trajectory is None: + _, driver_x_trajectory, driver_y_trajectory = self.driver_guiding_trajectory(num_steps=num_steps, dacc_gradient=0.0) + + if not stage_copy.is_upramp() and not stage_copy.is_downramp(): + if len(driver_x_trajectory) != num_steps or len(driver_y_trajectory) != num_steps: + raise ValueError('The length of driver_x_trajectory and driver_y_trajectory must be the same as num_steps.') + + # Initialise arrays + s_trajectory = np.full(num_steps, None, dtype=float) + x_trajectory = np.full(num_steps, None, dtype=float) + y_trajectory = np.full(num_steps, None, dtype=float) + px_trajectory = np.full(num_steps, None, dtype=float) + py_trajectory = np.full(num_steps, None, dtype=float) + pz_trajectory = np.full(num_steps, None, dtype=float) + + # Recursive call from the upramp + if stage_copy.upramp is not None: + upramp = stage_copy.convert_PlasmaRamp(stage_copy.upramp) + stage_copy.upramp = upramp + if self.external_focusing: + upramp.external_focusing_gradient = stage_copy.external_focusing_gradient + num_steps_upramp = stage_copy.num_steps_upramp + + s_trajectory_upramp, x_trajectory_upramp, y_trajectory_upramp, px_trajectory_upramp, py_trajectory_upramp, pz_trajectory_upramp, _, _ = upramp._estimate_beam_trajectory(s0, x0, y0, px0, py0, pz0, q, num_steps_upramp, driver_x_trajectory, driver_y_trajectory) + + # Initial parameters for the flattop + prop_length = s_trajectory_upramp[-1] + s_trajectory[:len(s_trajectory_upramp)] = s_trajectory_upramp + x0 = x_trajectory_upramp[-1] + x_trajectory[:len(x_trajectory_upramp)] = x_trajectory_upramp + y0 = y_trajectory_upramp[-1] + y_trajectory[:len(y_trajectory_upramp)] = y_trajectory_upramp + px0 = px_trajectory_upramp[-1] + px_trajectory[:len(px_trajectory_upramp)] = px_trajectory_upramp + py0 = py_trajectory_upramp[-1] + py_trajectory[:len(py_trajectory_upramp)] = py_trajectory_upramp + pz0 = pz_trajectory_upramp[-1] + pz_trajectory[:len(pz_trajectory_upramp)] = pz_trajectory_upramp + + i = num_steps_upramp - 1 + + if stage_copy.downramp is not None: + i_end = stage_copy.idx_flat_end + else: + i_end = num_steps - 1 + + # No ramps + else: + prop_length = s0 + s_trajectory[0] = prop_length + x_trajectory[0] = x0 + y_trajectory[0] = y0 + px_trajectory[0] = px0 + py_trajectory[0] = py0 + pz_trajectory[0] = pz0 + + i = 0 + i_end = num_steps - 1 + + if self.is_downramp(): + # Extract the part of drive beam trajectory for the downramp. Note one element longer than num_steps_downramp. + driver_x_trajectory = driver_x_trajectory[stage_copy.idx_flat_end:] + driver_y_trajectory = driver_y_trajectory[stage_copy.idx_flat_end:] + + # Set the initial conditions + x = x0 + y = y0 + px = px0 + py = py0 + pz = pz0 + + # Solve the equations of motion + while i < i_end: + + # Drift + prop_length = prop_length + 1/2*ds + x = x + px/pz*1/2*ds + y = y + py/pz*1/2*ds + + # Kick + dpx = q*g*x*ds - q*g0*driver_x_trajectory[i]*ds + px = px + dpx + dpy = q*g*y*ds - q*g0*driver_y_trajectory[i]*ds + py = py + dpy + dpz = - q * self.nom_accel_gradient_flattop * ds/SI.c + pz = pz + dpz + + # Drift + prop_length = prop_length + 1/2*ds + x = x + px/pz*1/2*ds + y = y + py/pz*1/2*ds + + i = i + 1 + s_trajectory[i] = prop_length + x_trajectory[i] = x + y_trajectory[i] = y + px_trajectory[i] = px + py_trajectory[i] = py + pz_trajectory[i] = pz + + # Recursive call from the downramp + if stage_copy.downramp is not None: + downramp = stage_copy.convert_PlasmaRamp(stage_copy.downramp) + stage_copy.downramp = downramp + if self.external_focusing: + downramp.external_focusing_gradient = stage_copy.external_focusing_gradient + + num_steps_downramp = stage_copy.num_steps_downramp + + s_trajectory_downramp, x_trajectory_downramp, y_trajectory_downramp, px_trajectory_downramp, py_trajectory_downramp, pz_trajectory_downramp, _, _ = downramp._estimate_beam_trajectory(prop_length, x, y, px, py, pz, q, num_steps_downramp+1, driver_x_trajectory, driver_y_trajectory) + + s_trajectory[-len(s_trajectory_downramp)+1:] = s_trajectory_downramp[1:] # Dsicarding the first element in s_trajectory_downramp + x_trajectory[-len(x_trajectory_downramp)+1:] = x_trajectory_downramp[1:] + y_trajectory[-len(y_trajectory_downramp)+1:] = y_trajectory_downramp[1:] + px_trajectory[-len(px_trajectory_downramp)+1:] = px_trajectory_downramp[1:] + py_trajectory[-len(py_trajectory_downramp)+1:] = py_trajectory_downramp[1:] + pz_trajectory[-len(pz_trajectory_downramp)+1:] = pz_trajectory_downramp[1:] + + return s_trajectory, x_trajectory, y_trajectory, px_trajectory, py_trajectory, pz_trajectory, driver_x_trajectory, driver_y_trajectory + # ================================================== # Apply waterfall function to all beam dump files def __waterfall_fcn(self, fcns, edges, data_dir, species='beam', remove_halo_nsigma=None, args=None): @@ -698,7 +1506,7 @@ def __waterfall_fcn(self, fcns, edges, data_dir, species='beam', remove_halo_nsi Applies waterfall function to all HiPACE++ HDF5 output files in ``data_dir``. - Parameters + Parameters ---------- fcns : A list of ``Beam`` class methods Beam class profile methods such as ``Beam.current_profile``, @@ -727,7 +1535,7 @@ def __waterfall_fcn(self, fcns, edges, data_dir, species='beam', remove_halo_nsi Returns - ---------- + ------- waterfalls : list of 2D float ndarrays Each element in ``waterfalls`` corresponds to the output of one function in ``fcns`` applied across all files (i.e., simulation @@ -803,7 +1611,7 @@ def plot_waterfalls(self, data_dir, species='beam', remove_halo_nsigma=20, save_ Returns - ---------- + ------- ``None`` ''' diff --git a/src/abel/classes/stage/impl/stage_reduced_models.py b/src/abel/classes/stage/impl/stage_reduced_models.py index ae36a750..995bfb8a 100644 --- a/src/abel/classes/stage/impl/stage_reduced_models.py +++ b/src/abel/classes/stage/impl/stage_reduced_models.py @@ -682,9 +682,9 @@ def track_upramp(self, beam0, driver0, shot_path=None): raise TypeError('upramp is not a StageReducedModels.') # Set a new time step for the ramp - n_steps = 25 # Do n_steps time steps in the tracking of the ramp + n_steps = 15 # Do n_steps time steps in the tracking of the ramp lambda_beta = upramp.matched_beta_function_flattop(beam0_energy) * 2*np.pi # [m], betatron wavelength - upramp.time_step_mod = min(upramp.length_flattop / (lambda_beta*n_steps), 0.02) # Step size in in units of betatron wavelength, equivalent to time step size in units of betatron wavelength/c. + upramp.time_step_mod = min(upramp.length_flattop / (lambda_beta*n_steps), 0.04) # Step size in in units of betatron wavelength, equivalent to time step size in units of betatron wavelength/c. elif type(self.upramp) is Stage: upramp = self.upramp # Allow for other types of ramps @@ -782,9 +782,9 @@ def track_downramp(self, beam0, driver0, shot_path=None): raise TypeError('downramp is not a StageReducedModels.') # Set a new time step for the ramp - n_steps = 25 # Do n_steps time steps in the tracking of the ramp + n_steps = 15 # Do n_steps time steps in the tracking of the ramp lambda_beta = downramp.matched_beta_function_flattop(beam0_energy) * 2*np.pi # [m], betatron wavelength - downramp.time_step_mod = min(downramp.length_flattop / (lambda_beta*n_steps), 0.02) # Step size in in units of betatron wavelength, equivalent to time step size in units of betatron wavelength/c. + downramp.time_step_mod = min(downramp.length_flattop / (lambda_beta*n_steps), 0.04) # Step size in in units of betatron wavelength, equivalent to time step size in units of betatron wavelength/c. elif type(self.downramp) is Stage: downramp = self.downramp # Allow for other types of ramps diff --git a/src/abel/classes/stage/stage.py b/src/abel/classes/stage/stage.py index 45b61dde..f1237d51 100644 --- a/src/abel/classes/stage/stage.py +++ b/src/abel/classes/stage/stage.py @@ -515,13 +515,39 @@ def _calc_ramp_length(self, ramp : Self) -> float: ramp_beta_mag = self.ramp_beta_mag else: raise ValueError('No ramp_beta_mag defined.') - - ramp_length = beta_matched(self.plasma_density, ramp.nom_energy)*np.pi/(2*np.sqrt(1/ramp_beta_mag)) + + g = SI.e*self.plasma_density/(2*SI.epsilon_0*SI.c) # [T/m], ion background focusing gradient + #if self.external_focusing_gradient is not None: # Add contribution from external field + # g = g + self.external_focusing_gradient # external_focusing_gradient may itself depend on the total length, so this may cause an infinite loop. + + k_beta = np.sqrt(g*SI.c/ramp.nom_energy) # [m^-1], betatron wavenumber. + ramp_length = 1/k_beta * np.pi/2 * np.sqrt(ramp_beta_mag) # k_beta*ramp_length = pi/2 gives pi/2 phase advance. + if ramp_length < 0.0: raise ValueError(f"ramp_length = {ramp_length} [m] < 0.0") return ramp_length + # ================================================== + def get_ramp_length(self) -> float: + """ + Get the length of the ramps if the stage has ramps already set up. + Returns 0.0 otherwise. + """ + + if self.has_ramp(): + if self.upramp is not None and self.upramp.length is not None: + upramp_length = self.upramp.length + else: + upramp_length = 0.0 + if self.downramp is not None and self.downramp.length is not None: + downramp_length = self.downramp.length + else: + downramp_length = 0.0 + return upramp_length + downramp_length + else: + return 0.0 + # ================================================== @property def parent(self) -> Self | None: @@ -1336,7 +1362,9 @@ def get_cost_breakdown(self): def matched_beta_function(self, energy_incoming, match_entrance=True): ''' Calculates the matched beta function of the stage. If there is an - upramp, the beta function is matched to the upramp by default. + upramp, the beta function the beta function is magnified by default so + that it shrinks to the correct size when it enters the main flattop + plasma stage. Parameters @@ -1385,8 +1413,343 @@ def matched_beta_function_flattop(self, energy): ''' return beta_matched(self.plasma_density, energy) + + + # ================================================== + def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, plasma_density=None, q=SI.e): + """ + Calculate the stage length that gives ``num_beta_osc`` betatron + oscillations for a particle with given initial energy ``initial_energy`` + in a uniform plasma stage (excluding ramps) with defined nominal + acceleration gradient and plasma density ``plasma_density``. + + Parameters + ---------- + num_beta_osc : float + Total number of design betatron oscillations that the electron + should perform through the plasma stage excluding ramps. + + initial_energy : [eV] float, optional + The initial energy of the particle at the start of the plasma stage. + Defaults to ``self.nom_energy``. + + plasma_density : [m^-3] float, optional + The plasma density of the plasma stage. Defaults to + ``self.plasma_density``. + + q : [C] float, optional + Particle charge. q * nom_accel_gradient must be positive. Defaults + to elementary charge. + + + Returns + ------- + length : [m] float + Length of the plasma stage excluding ramps matched to the given + number of betatron oscillations. + """ + + if initial_energy is None: + if self.nom_energy is None: + raise ValueError('Stage.nom_energy not set.') + initial_energy = self.nom_energy + initial_energy = initial_energy*SI.e # [J] + + if self.nom_accel_gradient_flattop is None: + raise ValueError('Stage.nom_accel_gradient_flattop not set.') + nom_accel_gradient = self.nom_accel_gradient_flattop + + if q * nom_accel_gradient < 0: + raise ValueError('q * nom_accel_gradient must be positive.') + + if plasma_density is None: + if self.plasma_density is None: + raise ValueError('Stage.plasma_density not set.') + plasma_density = self.plasma_density + + if num_beta_osc < 0: + raise ValueError('Number of input betatron oscillations must be positive.') + + g = SI.e*plasma_density/(2*SI.epsilon_0*SI.c) # [T/m] + gradient_prefactor = 2*np.sqrt(np.abs(q)*g*SI.c) / (q*nom_accel_gradient) + + length = ((2*np.pi*num_beta_osc/gradient_prefactor + np.sqrt(initial_energy))**2 - initial_energy) / (q*nom_accel_gradient) + + return length + + + # ================================================== + def calc_flattop_num_beta_osc(self, num_beta_osc): + """ + For a given total number of betatron oscillations ``num_beta_osc`` that + an electron with energy ``self.nom_energy`` will undergo across the + whole plasma stage including its uniform ramps, this function calculates + the number of betatron oscillations that should be performed in the main + flattop plasma stage by subtracting the contributions from the ramps. + + The contributions from the ramps are calculated from the phase advances + in the ramps as L_ramp/beta_matched_ramp. + + - If the stage does have ramps that have not been fully set up, a + deepcopy of the stage is created to set up its ramps using + :func:`Stage._prepare_ramps() `. + + - If the stage does not have ramps, will simply return the input total + number of betatron oscillations ``num_beta_osc``. + + + Parameters + ---------- + num_beta_osc : float + Total number of design betatron oscillations that the electron + should perform through the whole plasma stage including ramps. + + + Returns + ------- + float + The number of betatron oscillations the electron should undergo in + the flattop stage after the contributions from the ramps (if + applicable) have been subtracted from ``num_beta_osc``. + """ + + if num_beta_osc < 0: + raise ValueError('Number of input betatron oscillations must be positive.') + + if not self.has_ramp(): + return num_beta_osc + + if self.upramp.ramp_shape != 'uniform' or self.downramp.ramp_shape != 'uniform': + raise ValueError('This method assumes uniform ramps.') + + # Make a copy of the stage and set up its ramps if they are not set yp + ramps_not_set_up = ( + (self.upramp is not None and self.upramp.length is None) or + (self.downramp is not None and self.downramp.length is None) + ) + if ramps_not_set_up: + stage_copy = copy.deepcopy(self) + stage_copy._prepare_ramps() + else: + stage_copy = self + + # Calculate the upramp length and phase advance + if stage_copy.upramp is not None: + upramp_phase_advance = stage_copy.upramp.phase_advance_beta_evolution() + else: + upramp_phase_advance = 0.0 + + # Calculate the downramp length and matched beta function + if stage_copy.nom_energy_gain_flattop is None: + raise ValueError('Stage.nom_energy_gain_flattop not set.') + + # Calculate the downramp length and phase advance + if stage_copy.downramp is not None: + downramp_phase_advance = stage_copy.downramp.phase_advance_beta_evolution() + else: + downramp_phase_advance = 0.0 + + # Calculate the phase advance in the flattop stage + tot_phase_advance = num_beta_osc * 2*np.pi + flattop_phase_advance = tot_phase_advance - upramp_phase_advance - downramp_phase_advance + + return flattop_phase_advance/(2*np.pi) + + + # ================================================== + def length_flattop2num_beta_osc(self, length_flattop=None, initial_energy=None, nom_accel_gradient_flattop=None, plasma_density=None, q=SI.e): + """ + Calculate the number of betatron oscillations a particle can undergo in + the stage (excluding ramps). + + Will take into account the contribution from an external linear magnetic + field B=[gy,-gx,0] if :attr:`self.external_focusing_gradient ` + is not ``None``. + + Parameters + ---------- + length_flattop : [m] float, optional + Length of a plasma stage excluding ramps that the particle can + perform betatron scillations in. Defaults to + ``self.length_flattop``. + + initial_energy : [eV] float, optional + The initial energy of the particle at the start of the plasma stage. + Defaults to ``self.nom_energy``. + + nom_accel_gradient_flattop : [V/m] float, optional + Nominal accelerating gradient of the plasma stage exclusing ramps. + Defaults to ``self.nom_accel_gradient_flattop``. + + plasma_density : [m^-3] float, optional + The plasma density of the plasma stage. Defaults to + ``self.plasma_density``. + + q : [C] float, optional + Particle charge. q * nom_accel_gradient_flattop must be positive. + Defaults to elementary charge. + + + Returns + ------- + num_beta_osc : float + Total number of betatron oscillations that the particle will perform + across the plasma stage. + """ + + from abel.utilities.plasma_physics import k_p + + if length_flattop is None: + if self.length_flattop is None: + raise ValueError('Stage.length_flattop not set.') + length_flattop = self.length_flattop + + if initial_energy is None: + if self.nom_energy is None: + raise ValueError('Stage.nom_energy not set.') + initial_energy = self.nom_energy + + if nom_accel_gradient_flattop is None: + if self.nom_accel_gradient_flattop is None: + raise ValueError('Stage.nom_accel_gradient_flattop not set.') + nom_accel_gradient_flattop = self.nom_accel_gradient_flattop + + if q * nom_accel_gradient_flattop < 0: + raise ValueError('q * nom_accel_gradient_flattop must be positive.') + + if plasma_density is None: + plasma_density = self.plasma_density + + if nom_accel_gradient_flattop < 1e-15: # Need to treat very small gradients separately. Often the case for ramps. + return self.phase_advance_beta_evolution()/(2*np.pi) + else: + g = SI.e*plasma_density/(2*SI.epsilon_0*SI.c) # [T/m], ion background focusing gradient + if self.external_focusing_gradient is not None: + g = g + self.external_focusing_gradient + + prefactor = 2*np.sqrt(np.abs(q)*g*SI.c) / (q*nom_accel_gradient_flattop) + energy_scaling = np.sqrt(initial_energy*SI.e + q*nom_accel_gradient_flattop*length_flattop) - np.sqrt(initial_energy*SI.e) + + num_beta_osc = prefactor * energy_scaling / (2*np.pi) + + return num_beta_osc + + + # ================================================== + def match_length_2_num_beta_osc(self, num_beta_osc, q=SI.e): + """ + Set :attr:`self.length_flattop ` for a + uniform plasma stage such that a particle with initial energy + :attr:`self.nom_energy ` will perform + ``num_beta_osc`` betatron oscillations through the stage (including any + existing ramps). + + - Assumes that each of the (uniform) ramps are configured to give pi/2 + phase advance for the main beam. + + - The stage length calculation is performed using + :meth:`Stage.calc_flattop_num_beta_osc() `. + + + Parameters + ---------- + num_beta_osc : float + Total number of design betatron oscillations that the electron + should perform through the plasma stage excluding ramps. + + q : [C] float, optional + Particle charge. q * nom_accel_gradient must be positive. Defaults + to elementary charge. + + + Returns + ------- + None + """ + + # Assess whether length flattop can be set + if self._length_flattop_calc is not None and self._length_flattop is None: + from abel.classes.stage.stage import VariablesOverspecifiedError + raise VariablesOverspecifiedError("Stage length already known/calculateable, cannot set.") + + if self.has_ramp(): + # Calculate the number of betatron oscillations that the main beam + # should perform in the flattop: + if self.upramp.ramp_shape != 'uniform' or self.downramp.ramp_shape != 'uniform': + raise ValueError('This method assumes uniform ramps.') + if self.upramp.length_flattop is not None or self.downramp.length_flattop is not None: + raise ValueError('This method assumes uniform ramps with length set to give pi/2 phase advance for the main beam.') + + num_beta_osc_flattop = num_beta_osc - 0.5 # The ramps are by default set up to give pi/2 phase advance for the main beam. + # num_beta_osc_flattop = self.calc_flattop_num_beta_osc(num_beta_osc) + else: + num_beta_osc_flattop = num_beta_osc + + # Calculate the length of the flattop stage + length_flattop = self.calc_length_num_beta_osc(num_beta_osc=num_beta_osc_flattop, + initial_energy=self.nom_energy, + plasma_density=self.plasma_density, + q=q) + + # Set the length of the flattop stage + self.length_flattop = length_flattop + + + # ============================================= + @property + def external_focusing_gradient(self) -> float: + """ + Return None by default for ``Stage`` subclasses not supporting external + focusing fields. + """ + return None + # ================================================== + def phase_advance_beta_evolution(self, beta0=None): + """ + Calculate the phase advance in a stage by evolving the beta function + through a single element lattice set up using the stage's focusing + strength and length. The evolved beta function is then integrated along + the stage. + + Parameters + ---------- + beta0 : [m] float + The initial beta function at the start of the stage. If ``None``, + will calculate the matched beta function for the flattop stage and + scale it according the the ramp's ramp_beta_mag if ``self`` is a + ramp. + + + Returns + ------- + float + The phase advance calculated by integrating the beta function + evolution through the stage. + """ + + from abel.utilities.beam_physics import evolve_beta_function + from abel.utilities.beam_physics import phase_advance + + g = SI.e*self.plasma_density/(2*SI.epsilon_0) + if self.external_focusing_gradient is not None: + g = g + self.external_focusing_gradient * SI.c + p0 = np.sqrt((self.nom_energy*SI.e)**2-(SI.m_e*SI.c**2)**2)/SI.c + if beta0 is None: + if self.is_upramp(): + beta0 = self.matched_beta_function(self.nom_energy, match_entrance=True) + elif self.is_downramp(): + beta0 = beta_matched(self.parent.plasma_density, self.nom_energy) + else: + beta0 = beta_matched(self.plasma_density, self.nom_energy) + + ls = np.array([self.length_flattop]) + ks = np.array([g*SI.e/SI.c/p0]) + _, _, beta_evolution = evolve_beta_function(ls=ls, ks=ks, beta0=beta0, fast=False, plot=False) + return phase_advance(beta_evolution[0,:], beta_evolution[1,:]) + + # ================================================== def energy_usage(self): return self.driver_source.energy_usage() diff --git a/src/abel/utilities/beam_physics.py b/src/abel/utilities/beam_physics.py index 43ac3a82..2f19f680 100644 --- a/src/abel/utilities/beam_physics.py +++ b/src/abel/utilities/beam_physics.py @@ -1152,3 +1152,78 @@ def evolve_chromatic_amplitude(ls, inv_rhos, ks, ms, taus, beta0, alpha0=0, Dx0= return W, evolution + +# ============================================= +def phase_advance(ss, betas): + """ + Calculate the phase advance in one dimesion by using the composite Simpson’s + rule (:func:`scipy.integrate.simpson() `) to + integrate two arrays containing the location and the beta function. + """ + + from scipy import integrate + inv_betas = 1/betas + return integrate.simpson(y=inv_betas, x=ss) + + +# ============================================= +def phase_advance_traj_data(pos, angles, orbit_pos=None, orbit_angles=None): + """ + Compute the unwrapped betatron phase from transverse position and slope + data. + + Parameters + ---------- + pos : [m] 1D float ndarray + Transverse position along the beamline. + + angles : 1D float ndarray + Transverse angle dpos/ds. + + orbit_pos : [m] 1D float ndarray, optional + Reference orbit transverse position (e.g., driver trajectory). If + provided, these are subtracted before phase calculation. Defaults to + `None`. + + orbit_angles : 1D float ndarray, optional + Reference orbit transverse angles (e.g., driver trajectory angles). If + provided, these are subtracted before phase calculation. Defaults to + `None`. + + Returns + ------- + phase_advance : float + Unwrapped phase angle in radians, representing the continuous betatron + phase advance at the end of the beamline + """ + + if orbit_pos is not None: + pos = pos - orbit_pos + if orbit_angles is not None: + angles = angles - orbit_angles + + # Compute the phase space angle by normalising coordinates to a unit circle + # and taking the arctangent before unwrapping + phase = np.unwrap(np.arctan2(angles / np.max(np.abs(angles)), + pos / np.max(np.abs(pos)))) + + return np.abs(phase[-1] - phase[0]) + + +# ============================================= +def arc_lengths(s_trajectory, x_trajectory): + """ + Docstring for arc_length + + :param s_trajectory: Description + :param x_trajectory: Description + """ + + ds = np.diff(s_trajectory) + dx = np.diff(x_trajectory) + + length = np.cumsum(np.sqrt(ds**2 + dx**2)) + + length = np.insert(length, 0, 0.0) + + return length diff --git a/src/abel/wrappers/hipace/hipace_wrapper.py b/src/abel/wrappers/hipace/hipace_wrapper.py index 7a0b9177..90ce9487 100644 --- a/src/abel/wrappers/hipace/hipace_wrapper.py +++ b/src/abel/wrappers/hipace/hipace_wrapper.py @@ -156,7 +156,9 @@ def hipace_write_inputs(filename_input, filename_beam, filename_driver, plasma_d print('>> HiPACE++: Changing from', num_cell_xy, 'to', new_num_cell_xy, ' (i.e., 2^n-1) for better performance.') num_cell_xy = new_num_cell_xy - # plasma-density profile from file + # gradient for external magnetic field + if external_focusing_gradient is None: + external_focusing_gradient = '#' if abs(external_focusing_gradient) > 0: external_focusing_comment = '' else: diff --git a/tests/test_StageHipace.py b/tests/test_StageHipace.py new file mode 100644 index 00000000..d0389c38 --- /dev/null +++ b/tests/test_StageHipace.py @@ -0,0 +1,370 @@ +# -*- coding: utf-8 -*- +# This file is part of ABEL +# Copyright 2025, The ABEL Authors +# Authors: C.A.Lindstrøm(1), J.B.B.Chen(1), O.G.Finnerud(1), D.Kalvik(1), E.Hørlyk(1), A.Huebl(2), K.N.Sjobak(1), E.Adli(1) +# Affiliations: 1) University of Oslo, 2) LBNL +# License: GPL-3.0-or-later + + +""" +ABEL : StageHipace unit tests +""" + +import pytest +from abel import * +#import shutil +import numpy as np + + +def setup_trapezoid_driver_source(enable_xy_jitter=False, enable_xpyp_jitter=False, x_angle=0.0, y_angle=0.0): + driver = SourceTrapezoid() + driver.current_head = 0.1e3 # [A] + driver.bunch_length = 1050e-6 # [m] + driver.z_offset = 1602e-6 # [m] + driver.x_angle = x_angle # [rad] + driver.y_angle = y_angle # [rad] + + driver.num_particles = 1000000 + driver.charge = 5.0e10 * -SI.e # [C] + driver.energy = 50e9 # [eV] + driver.gaussian_blur = 50e-6 # [m] + driver.rel_energy_spread = 0.01 + + driver.emit_nx = 50e-6 * driver.energy/4.5e9 # [m rad] + driver.emit_ny = 100e-6 * driver.energy/4.5e9 # [m rad] + driver.beta_x, driver.beta_y = 0.5, 0.5 # [m] + + if enable_xy_jitter: + driver.jitter.x = 100e-9 # [m], std + driver.jitter.y = 100e-9 # [m], std + + if enable_xpyp_jitter: + driver.jitter.xp = 1.0e-6 # [rad], std + driver.jitter.yp = 1.0e-6 # [rad], std + + driver.symmetrize = True + + return driver + + +def setup_minimal_StageHipace(nom_energy=100e9, plasma_density=6e20, external_focusing=False, nom_accel_gradient_flattop=1e9): + stage = StageHipace() + stage.nom_energy = nom_energy # [eV] + stage.plasma_density = plasma_density # [m^-3] + stage.driver_source = setup_trapezoid_driver_source() + stage.external_focusing = external_focusing + stage.nom_accel_gradient_flattop = nom_accel_gradient_flattop # [V/m] + + return stage + + +@pytest.mark.StageHipace +def test_external_focusing(): + """ + Tests for ``StageHipace.external_focusing_gradient`` for setting and + calculating the gradient g_ext [T/m] for an external azimuthal magnetic + field B = [g_ext*y, -g_ext*x, 0]. + """ + + # ========== Tests withexternal focusing disabled ========== + stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=False, nom_accel_gradient_flattop=None) + assert stage.external_focusing is False + assert stage._external_focusing_gradient is None + assert stage.external_focusing_gradient is None + with pytest.raises(ValueError): + stage.external_focusing_gradient = 3.14 # Cannot set any value when stage.external_focusing is False. + + + # ========== Tests with external focusing enabled ========== + stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True, nom_accel_gradient_flattop=None) + assert stage.external_focusing is True + assert stage._external_focusing_gradient is None + assert stage.external_focusing_gradient is None + with pytest.raises(ValueError): + stage.external_focusing_gradient = 'test' + stage.external_focusing_gradient = 3.14 + assert np.isclose(3.14, stage.external_focusing_gradient, rtol=1e-10, atol=0.0) + stage.external_focusing_gradient = 314.0 + assert np.isclose(314.0, stage.external_focusing_gradient, rtol=1e-10, atol=0.0) + + + # ========== Test the external gradient calculation without ramps ========== + stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True, nom_accel_gradient_flattop=1e9) + stage.driver_source.energy = 4.5e9 # [eV] + assert stage.external_focusing is True + assert stage._external_focusing_gradient is None + assert stage.external_focusing_gradient is None + assert np.isclose(1.0, stage.driver_half_oscillations, rtol=1e-10, atol=0.0) + stage.length = 1.2945695557961696 # [m] + assert np.isclose(1.2945695557961696*1e9, stage.nom_energy_gain, rtol=1e-10, atol=0.0) + assert np.isclose(1.2945695557961696*1e9, stage.nom_energy_gain_flattop, rtol=1e-10, atol=0.0) + assert np.isclose(88.3976616856249, stage.external_focusing_gradient, rtol=1e-10, atol=0.0) + + stage.external_focusing_gradient = 3.14 # Overwrite the focusing gradient + assert np.isclose(3.14, stage.external_focusing_gradient, rtol=1e-10, atol=0.0) + stage.external_focusing_gradient = None # Overwrite the focusing gradient + assert np.isclose(88.3976616856249, stage.external_focusing_gradient, rtol=1e-10, atol=0.0) # Should again automatically calculate a focusing gradient. + + # Match the focusing gradient to more driver betatron oscillations + stage.driver_half_oscillations = 2.0 + stage.length = 1.2841918239865389 + assert np.isclose(88.3976616856249, stage.external_focusing_gradient, rtol=1e-10, atol=0.0) # Still the previous focusing gradient. + stage.external_focusing_gradient = None # Overwrite the focusing gradient + assert np.isclose(359.3285677857948, stage.external_focusing_gradient, rtol=1e-10, atol=0.0) # Should again automatically calculate a focusing gradient. + + + # ========== Test the external gradient calculation with ramps ========== + stage = setup_minimal_StageHipace(nom_energy=10e9, external_focusing=True, nom_accel_gradient_flattop=1e9) + stage.driver_source.energy = 20e9 # [eV] + stage.ramp_beta_mag = 2.0 + stage.upramp = PlasmaRamp() + stage.downramp = PlasmaRamp() + assert stage.external_focusing is True + assert stage._external_focusing_gradient is None + assert stage.external_focusing_gradient is None + assert np.isclose(1.0, stage.driver_half_oscillations, rtol=1e-10, atol=0.0) + stage.length_flattop = 1.3834379291754226 # [m] + + assert stage._external_focusing_gradient is None + assert np.isclose(263.5820977342094, stage.external_focusing_gradient, rtol=1e-5, atol=0.0) + assert np.isclose(263.5820977342094, stage._external_focusing_gradient, rtol=1e-5, atol=0.0) + + stage.external_focusing_gradient = 3.14 # Overwrite the focusing gradient + assert np.isclose(3.14, stage.external_focusing_gradient, rtol=1e-10, atol=0.0) + stage.external_focusing_gradient = None # Overwrite the focusing gradient + assert np.isclose(263.5820977342094, stage.external_focusing_gradient, rtol=1e-10, atol=0.0) # Should again automatically calculate a focusing gradient. + + # Match the focusing gradient to more driver betatron oscillations + stage.driver_half_oscillations = 2.0 + stage.length_flattop = 1.3520401015419283 # [m] + assert np.isclose(2.0, stage.driver_half_oscillations, rtol=1e-10, atol=0.0) + assert np.isclose(263.5820977342094, stage.external_focusing_gradient, rtol=1e-10, atol=0.0) # Still the previous focusing gradient. + stage.external_focusing_gradient = None # Overwrite the focusing gradient + assert np.isclose(1097.699359490811, stage.external_focusing_gradient, rtol=1e-5, atol=0.0) # Should again automatically calculate a focusing gradient. + assert np.isclose(1097.699359490811, stage._external_focusing_gradient, rtol=1e-5, atol=0.0) + + +@pytest.mark.StageHipace +def test_calc_length_num_beta_osc(): + """ + Tests for ``StageHipace.calc_length_num_beta_osc()`` for accurately matching the + stage length and external driver guiding field gradient to a desired number + of drive beam and main beam betatron oscillations. + """ + + num_beta_osc = 4.0 # The number of betatron oscilations that the main beam is expected to perform. + + stage = setup_minimal_StageHipace() + + # ========== Tests without any external fields ========== + assert stage.external_focusing is False + assert stage._external_focusing_gradient is None + + stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc, initial_energy=stage.nom_energy, driver_half_oscillations=1.0) + + assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-5, atol=0.0) + assert stage._external_focusing_gradient is None + assert np.isclose(3.440220555221998, stage.length_flattop, rtol=1e-5, atol=0.0) + + + stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc, driver_half_oscillations=2.0) + + assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-5, atol=0.0) + assert stage._external_focusing_gradient is None + assert np.isclose(3.440220555221998, stage.length_flattop, rtol=1e-5, atol=0.0) + + + # ========== Tests with external fields ========== + stage = setup_minimal_StageHipace(external_focusing=True) + stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc, driver_half_oscillations=1.0) + assert stage._external_focusing_gradient is None + stage._external_focusing_gradient = 140.1695315279373 + + assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0) + assert np.isclose(3.4268706541720553, stage.length_flattop, rtol=1e-5, atol=0.0) + + + stage = setup_minimal_StageHipace(external_focusing=True) + stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc, driver_half_oscillations=2.0) + assert stage._external_focusing_gradient is None + stage._external_focusing_gradient = 574.1247168375805 + + assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0) + assert np.isclose(3.3865024719148242, stage.length_flattop, rtol=1e-5, atol=0.0) + + + # ========== With external fields, lower stage nominal energy ========== + num_beta_osc2 = 8.0 + stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True) + stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc2, driver_half_oscillations=1.0) + assert stage._external_focusing_gradient is None + stage._external_focusing_gradient = 1038.1202586404845 + + assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0) + assert np.isclose(1.259217324849329, stage.length_flattop, rtol=1e-5, atol=0.0) + + + stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True) + stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc2, driver_half_oscillations=2.0) + assert stage._external_focusing_gradient is None + stage._external_focusing_gradient = 5119.8513420067175 + + assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0) + assert np.isclose(1.13403339406399, stage.length_flattop, rtol=1e-5, atol=0.0) + + + # ========== With external fields, lower stage nominal energy, lower driver energy ========== + stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True) + stage.driver_source.energy = 4.5e9 # [eV] + stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc2, driver_half_oscillations=1.0) + assert stage._external_focusing_gradient is None + stage._external_focusing_gradient = 88.3976616856249 + + assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0) + assert np.isclose(1.2945695557961696, stage.length_flattop, rtol=1e-5, atol=0.0) + + + stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True) + stage.driver_source.energy = 4.5e9 # [eV] + stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc2, driver_half_oscillations=2.0) + assert stage._external_focusing_gradient is None + stage._external_focusing_gradient = 359.3285677857948 + + assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0) + assert np.isclose(1.2841918239865389, stage.length_flattop, rtol=1e-5, atol=0.0) + + + +@pytest.mark.StageHipace +def test_match_length_2_num_beta_osc(): + """ + Tests for ``StageHipace.match_length_2_num_beta_osc()`` for + accurately matching the stage length and external driver guiding field + gradient to a desired number of drive beam and main beam betatron + oscillations. + """ + + num_beta_osc = 4.0 # The number of betatron oscilations that the main beam is expected to perform. + stage = setup_minimal_StageHipace() + + # ========== Tests without any external fields ========== + assert stage.external_focusing is False + assert stage._external_focusing_gradient is None + assert stage.external_focusing_gradient is None + assert np.isclose(1.0, stage.driver_half_oscillations, rtol=1e-15, atol=0.0) + + stage.match_length_2_num_beta_osc(num_beta_osc=num_beta_osc, driver_half_oscillations=1.0, set_consistent_params=True) + assert np.isclose(1.0, stage.driver_half_oscillations, rtol=1e-15, atol=0.0) + assert stage._external_focusing_gradient is None + assert stage.external_focusing_gradient is None + assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-5, atol=0.0) + assert np.isclose(3.440220555221998, stage.length_flattop, rtol=1e-5, atol=0.0) + + + stage.match_length_2_num_beta_osc(num_beta_osc=num_beta_osc, driver_half_oscillations=2.0, set_consistent_params=True) + assert np.isclose(1.0, stage.driver_half_oscillations, rtol=1e-15, atol=0.0) # Should still be the default value 1.0, since stage.external_focusing is False. + assert stage.external_focusing_gradient is None + assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-5, atol=0.0) + assert stage._external_focusing_gradient is None + assert np.isclose(3.440220555221998, stage.length_flattop, rtol=1e-5, atol=0.0) + + + # ========== Tests with external fields ========== + stage = setup_minimal_StageHipace(external_focusing=True) + assert stage._external_focusing_gradient is None + assert stage.external_focusing_gradient is None + stage.external_focusing_gradient = 3.14 # Set a random gradient which will be overwritten later + assert np.isclose(3.14, stage.external_focusing_gradient, rtol=1e-5, atol=0.0) + + stage.match_length_2_num_beta_osc(num_beta_osc=num_beta_osc) + + assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0) + assert np.isclose(140.1695315279373, stage._external_focusing_gradient, rtol=1e-5, atol=0.0) + assert np.isclose(140.1695315279373, stage.external_focusing_gradient, rtol=1e-5, atol=0.0) + assert np.isclose(3.4268706541720553, stage.length_flattop, rtol=1e-5, atol=0.0) + + + stage = setup_minimal_StageHipace(external_focusing=True) + stage.match_length_2_num_beta_osc(num_beta_osc=num_beta_osc, driver_half_oscillations=2.0) + + assert np.isclose(2.0, stage.driver_half_oscillations, rtol=1e-15, atol=0.0) + assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0) + assert np.isclose(574.1247168375805, stage._external_focusing_gradient, rtol=1e-5, atol=0.0) + assert np.isclose(574.1247168375805, stage.external_focusing_gradient, rtol=1e-5, atol=0.0) + assert np.isclose(3.3865024719148242, stage.length_flattop, rtol=1e-5, atol=0.0) + + + # ========== With external fields, lower stage nominal energy ========== + num_beta_osc2 = 8.0 + stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True) + stage.match_length_2_num_beta_osc(num_beta_osc=num_beta_osc2) + + assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0) + assert np.isclose(1038.1202586404845, stage._external_focusing_gradient, rtol=1e-5, atol=0.0) + assert np.isclose(1038.1202586404845, stage.external_focusing_gradient, rtol=1e-5, atol=0.0) + assert np.isclose(1.259217324849329, stage.length_flattop, rtol=1e-5, atol=0.0) + + + stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True) + stage.match_length_2_num_beta_osc(num_beta_osc=num_beta_osc2, driver_half_oscillations=2.0) + + assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0) + assert np.isclose(5119.8513420067175, stage._external_focusing_gradient, rtol=1e-5, atol=0.0) + assert np.isclose(5119.8513420067175, stage.external_focusing_gradient, rtol=1e-5, atol=0.0) + assert np.isclose(1.13403339406399, stage.length_flattop, rtol=1e-5, atol=0.0) + + + # ========== With external fields, lower stage nominal energy, lower driver energy ========== + stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True) + stage.driver_source.energy = 4.5e9 # [eV] + stage.match_length_2_num_beta_osc(num_beta_osc=num_beta_osc2) + + assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0) + assert np.isclose(88.3976616856249, stage._external_focusing_gradient, rtol=1e-5, atol=0.0) + assert np.isclose(88.3976616856249, stage.external_focusing_gradient, rtol=1e-5, atol=0.0) + assert np.isclose(1.2945695557961696, stage.length_flattop, rtol=1e-5, atol=0.0) + + + stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True) + stage.driver_source.energy = 4.5e9 # [eV] + stage.match_length_2_num_beta_osc(num_beta_osc=num_beta_osc2, driver_half_oscillations=2.0) + + assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0) + assert np.isclose(359.3285677857948, stage._external_focusing_gradient, rtol=1e-5, atol=0.0) + assert np.isclose(359.3285677857948, stage.external_focusing_gradient, rtol=1e-5, atol=0.0) + assert np.isclose(1.2841918239865389, stage.length_flattop, rtol=1e-5, atol=0.0) + + + # ========== Tests with external fields, with ramps ========== + num_beta_osc3 = 5.5 + stage = setup_minimal_StageHipace(nom_energy=10e9, external_focusing=True) + stage.driver_source.energy = 20e9 # [eV] + stage.ramp_beta_mag = 2.0 + stage.upramp = PlasmaRamp() + stage.downramp = PlasmaRamp() + + stage.match_length_2_num_beta_osc(num_beta_osc=num_beta_osc3, driver_half_oscillations=1.0) + + assert np.isclose(263.5820977342094, stage._external_focusing_gradient, rtol=1e-5, atol=0.0) + assert np.isclose(263.5820977342094, stage.external_focusing_gradient, rtol=1e-5, atol=0.0) + assert np.isclose(1.3834379291754226, stage.length_flattop, rtol=1e-5, atol=0.0) + assert np.isclose(num_beta_osc3 - 0.5, stage.length_flattop2num_beta_osc(), rtol=1e-4, atol=0.0) + + + stage = setup_minimal_StageHipace(nom_energy=10e9, external_focusing=True) + stage.driver_source.energy = 20e9 # [eV] + stage.ramp_beta_mag = 2.0 + stage.upramp = PlasmaRamp() + stage.downramp = PlasmaRamp() + + stage.match_length_2_num_beta_osc(num_beta_osc=num_beta_osc3, driver_half_oscillations=2.0) + + assert np.isclose(2.0, stage.driver_half_oscillations, rtol=1e-10, atol=0.0) + assert np.isclose(1097.699359490811, stage._external_focusing_gradient, rtol=1e-5, atol=0.0) + assert np.isclose(1097.699359490811, stage.external_focusing_gradient, rtol=1e-5, atol=0.0) + assert np.isclose(1.3520401015419283, stage.length_flattop, rtol=1e-5, atol=0.0) + assert np.isclose(num_beta_osc3 - 0.5, stage.length_flattop2num_beta_osc(), rtol=1e-4, atol=0.0) + + + + diff --git a/tests/test_StageReducedModels_beamline.py b/tests/test_StageReducedModels_beamline.py index e9fccef4..9978dc2c 100644 --- a/tests/test_StageReducedModels_beamline.py +++ b/tests/test_StageReducedModels_beamline.py @@ -12,7 +12,7 @@ import pytest from abel import * import matplotlib -import shutil +import shutil, copy matplotlib.use('Agg') # Use a backend that does not display figure to suppress plots. def setup_trapezoid_driver_source(enable_xy_jitter=False, enable_xpyp_jitter=False): @@ -72,15 +72,15 @@ def setup_basic_main_source(plasma_density, ramp_beta_mag=1.0, energy=361.8e9): return main -def setup_StageReducedModels(plasma_density, driver_source, main_source, ramp_beta_mag, enable_tr_instability=True, enable_radiation_reaction=True, enable_ion_motion=False, use_ramps=False, drive_beam_update_period=0, save_final_step=False): +def setup_StageReducedModels(plasma_density, driver_source, ramp_beta_mag, length_flattop=1.56, enable_tr_instability=True, enable_radiation_reaction=True, enable_ion_motion=False, use_ramps=False, drive_beam_update_period=0, save_final_step=False): stage = StageReducedModels() - stage.time_step_mod = 0.03*2 # In units of betatron wavelengths/c. - stage.nom_energy_gain = 7.8e9/5 # [eV] - stage.length_flattop = 7.8/5 # [m] + stage.time_step_mod = 0.03*2 # In units of betatron wavelengths/c. + stage.length_flattop = length_flattop # [m] + if length_flattop is not None: + stage.nom_energy_gain = stage.length_flattop*1e9 # [eV] stage.plasma_density = plasma_density # [m^-3] stage.driver_source = driver_source - stage.main_source = main_source stage.ramp_beta_mag = ramp_beta_mag stage.enable_tr_instability = enable_tr_instability stage.enable_radiation_reaction = enable_radiation_reaction @@ -136,11 +136,10 @@ def test_baseline_linac(): enable_radiation_reaction = False enable_ion_motion = False use_ramps = False - drive_beam_update_period = 0 driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter) main_source = setup_basic_main_source(plasma_density, ramp_beta_mag=1.0) - stage = setup_StageReducedModels(plasma_density, driver_source, main_source, 1.0, enable_tr_instability, enable_radiation_reaction, enable_ion_motion, use_ramps, drive_beam_update_period) + stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, ramp_beta_mag=1.0, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps) interstage = setup_InterstageImpactX(stage) linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False) @@ -211,6 +210,7 @@ def test_ramped_linac(): """ np.random.seed(42) + from abel.utilities.plasma_physics import beta_matched num_stages = 2 plasma_density = 6.0e+20 # [m^-3] @@ -221,14 +221,28 @@ def test_ramped_linac(): enable_radiation_reaction = False enable_ion_motion = False use_ramps = True - drive_beam_update_period = 0 driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter) - main_source = setup_basic_main_source(plasma_density, ramp_beta_mag) - stage = setup_StageReducedModels(plasma_density, driver_source, main_source, ramp_beta_mag, enable_tr_instability, enable_radiation_reaction, enable_ion_motion, use_ramps, drive_beam_update_period) + driver_source.z_offset = 1602e-6 # [m] + main_source = setup_basic_main_source(plasma_density, ramp_beta_mag, energy=3.0e9) + + stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, ramp_beta_mag=ramp_beta_mag, length_flattop=None, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps) + + stage.nom_energy = main_source.energy + stage.nom_accel_gradient_flattop = 1e9 # [V/m] + + # Adjust the lengths of the two stages in the linac to match the number of betatron oscillation + stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc=9.5) + + last_stage = copy.deepcopy(stage) + last_stage.length_flattop = last_stage.calc_length_num_beta_osc(num_beta_osc=9.5, initial_energy=main_source.energy+stage.nom_energy_gain_flattop) + last_stage.nom_energy_gain = last_stage.length_flattop*1e9 + + # Set up the interstage interstage = setup_InterstageImpactX(stage) - linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False) + # Assemble the linac + linac = PlasmaLinac(source=main_source, stage=stage, last_stage=last_stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False) # Perform tracking linac.run('test_ramped_linac', overwrite=True, verbose=False) @@ -239,15 +253,16 @@ def test_ramped_linac(): assert linac.stage.driver_source.align_beam_axis is True assert stages[0].driver_source.align_beam_axis is True assert stages[1].driver_source.align_beam_axis is True - assert np.isclose(linac.nom_energy, 364920000000.0, rtol=1e-5, atol=0.0) + assert np.isclose(linac.nom_energy, 6462752491.36345, rtol=1e-5, atol=0.0) assert np.isclose(linac.nom_energy, stages[-1].nom_energy + stages[-1].nom_energy_gain, rtol=1e-5, atol=0.0) - assert np.isclose(linac.get_length(), 41.586517149595515, rtol=1e-5, atol=0.0) + assert np.isclose(linac.nom_energy, main_source.energy + stages[0].nom_energy_gain + stages[-1].nom_energy_gain, rtol=1e-5, atol=0.0) + assert np.isclose(linac.get_length(), 7.775398161868525, rtol=1e-5, atol=0.0) - assert np.isclose(stages[0].nom_energy, 361.8e9) - assert np.isclose(stages[0].nom_energy_gain, 1560000000.0, rtol=1e-15, atol=0.0) - assert np.isclose(stages[1].nom_energy, stages[0].nom_energy + stages[0].nom_energy_gain) - assert np.isclose(stages[1].nom_energy_gain, 1560000000.0, rtol=1e-15, atol=0.0) - assert np.isclose(stages[-1].upramp.length, 0.5747274948097144, rtol=1e-2, atol=0.0) + assert np.isclose(stages[0].nom_energy, 3.0e9) + assert np.isclose(stages[0].nom_energy_gain, 1567293075.7156029, rtol=1e-15, atol=0.0) + assert np.isclose(stages[1].nom_energy, stages[0].nom_energy + stages[0].nom_energy_gain, rtol=1e-15, atol=0.0) + assert np.isclose(stages[1].nom_energy_gain, 1895459415.6478472, rtol=1e-15, atol=0.0) + assert np.isclose(stages[-1].upramp.length, 0.06443515179509757, rtol=1e-2, atol=0.0) assert np.isclose(stages[-1].upramp.length_flattop, stages[-1].upramp.length, rtol=1e-15, atol=0.0) assert np.isclose(stages[-1].upramp.nom_energy, stages[-1].nom_energy, rtol=1e-15, atol=0.0) assert np.isclose(stages[-1].upramp.nom_energy_flattop, stages[-1].upramp.nom_energy, rtol=1e-15, atol=0.0) @@ -255,7 +270,7 @@ def test_ramped_linac(): assert np.isclose(stages[-1].upramp.nom_energy_gain_flattop, stages[-1].upramp.nom_energy_gain, rtol=1e-15, atol=0.0) assert np.isclose(stages[-1].upramp.nom_accel_gradient, 0.0, rtol=1e-15, atol=0.0) assert np.isclose(stages[-1].upramp.nom_accel_gradient_flattop, stages[-1].upramp.nom_accel_gradient, rtol=1e-15, atol=0.0) - assert np.isclose(stages[-1].downramp.length, 0.5759599015745918, rtol=1e-2, atol=0.0) + assert np.isclose(stages[-1].downramp.length, 0.07664823833842702, rtol=1e-2, atol=0.0) assert np.isclose(stages[-1].downramp.length_flattop, stages[-1].downramp.length, rtol=1e-15, atol=0.0) assert np.isclose(stages[-1].downramp.nom_energy, stages[-1].nom_energy + stages[-1].nom_energy_gain, rtol=1e-15, atol=0.0) assert np.isclose(stages[-1].downramp.nom_energy_flattop, stages[-1].downramp.nom_energy, rtol=1e-15, atol=0.0) @@ -284,7 +299,7 @@ def test_ramped_linac(): assert np.isclose(final_beam.energy(), stages[-1].nom_energy + stages[-1].nom_energy_gain, rtol=1e-2, atol=0.0) assert np.isclose(final_beam.bunch_length(), main_source.bunch_length, rtol=1e-1, atol=0.0) assert np.isclose(final_beam.charge(), main_source.charge, rtol=1e-5, atol=0.0) - assert np.isclose(final_beam.rel_energy_spread(), 0.02071572458689477, rtol=1e-1, atol=0.0) + assert np.isclose(final_beam.rel_energy_spread(), 0.01815229983882398, rtol=1e-1, atol=0.0) nom_beam_size_x = (stages[0].nom_energy/stages[-1].nom_energy)**(1/4)*initial_beam.beam_size_x() nom_beam_size_y = (stages[0].nom_energy/stages[-1].nom_energy)**(1/4)*initial_beam.beam_size_y() @@ -297,6 +312,40 @@ def test_ramped_linac(): assert np.isclose(final_beam.norm_emittance_x(), main_source.emit_nx, rtol=1e-1, atol=0.0) assert np.isclose(final_beam.norm_emittance_y(), main_source.emit_ny, rtol=1e-1, atol=0.0) + # Check ramp parameters + upramp1 = stages[0].upramp + downramp1 = stages[0].downramp + assert np.isclose(upramp1.plasma_density, stage.plasma_density/upramp1.ramp_beta_mag, rtol=1e-10, atol=0.0) + assert np.isclose(downramp1.plasma_density, stage.plasma_density/downramp1.ramp_beta_mag, rtol=1e-10, atol=0.0) + upramp1_length = beta_matched(stage.plasma_density, upramp1.nom_energy) * np.pi/2 * np.sqrt(upramp1.ramp_beta_mag) + assert np.isclose(upramp1.length_flattop, upramp1_length, rtol=1e-10, atol=0.0) + downramp1_length = beta_matched(stage.plasma_density, downramp1.nom_energy) * np.pi/2 * np.sqrt(downramp1.ramp_beta_mag) + assert np.isclose(downramp1.length_flattop, downramp1_length, rtol=1e-10, atol=0.0) + upramp2 = stages[1].upramp + downramp2 = stages[1].downramp + assert np.isclose(upramp2.plasma_density, stage.plasma_density/upramp2.ramp_beta_mag, rtol=1e-10, atol=0.0) + assert np.isclose(downramp2.plasma_density, stage.plasma_density/downramp2.ramp_beta_mag, rtol=1e-10, atol=0.0) + upramp2_length = beta_matched(stage.plasma_density, upramp2.nom_energy) * np.pi/2 * np.sqrt(upramp2.ramp_beta_mag) + assert np.isclose(upramp2.length_flattop, upramp2_length, rtol=1e-10, atol=0.0) + downramp2_length = beta_matched(stage.plasma_density, downramp2.nom_energy) * np.pi/2 * np.sqrt(downramp2.ramp_beta_mag) + assert np.isclose(downramp2.length_flattop, downramp2_length, rtol=1e-10, atol=0.0) + + # Check phase advances in ramps and flattop stage + assert np.isclose(upramp1.phase_advance_beta_evolution(), np.pi/2, rtol=1e-3, atol=0.0) # Calculate the phase advance by integrating the beta function. + assert np.isclose(upramp1.length_flattop2num_beta_osc(), 1/4, rtol=1e-3, atol=0.0) # Calculate the number of betatron oscillations a particle can undergo along the ramp. + assert np.isclose(downramp1.phase_advance_beta_evolution(), np.pi/2, rtol=1e-3, atol=0.0) + assert np.isclose(downramp1.length_flattop2num_beta_osc(), 1/4, rtol=1e-3, atol=0.0) + + assert np.isclose(upramp2.phase_advance_beta_evolution(), np.pi/2, rtol=1e-3, atol=0.0) # Calculate the phase advance by integrating the beta function. + assert np.isclose(upramp2.length_flattop2num_beta_osc(), 1/4, rtol=1e-3, atol=0.0) # Calculate the number of betatron oscillations a particle can undergo along the ramp. + assert np.isclose(downramp2.phase_advance_beta_evolution(), np.pi/2, rtol=1e-3, atol=0.0) + assert np.isclose(downramp2.length_flattop2num_beta_osc(), 1/4, rtol=1e-3, atol=0.0) + + assert np.isclose(stages[0].length_flattop2num_beta_osc(), 9.5, rtol=1e-3, atol=0.0) # Calculate the number of betatron oscillations a particle can undergo along the stage. + assert np.isclose(stages[0].length_flattop2num_beta_osc() + upramp1.length_flattop2num_beta_osc() + downramp1.length_flattop2num_beta_osc(), 10.0, rtol=1e-3, atol=0.0) # Calculate the total phase adcance. + assert np.isclose(stages[1].length_flattop2num_beta_osc(), 9.5, rtol=1e-3, atol=0.0) # Calculate the number of betatron oscillations a particle can undergo along the stage. + assert np.isclose(stages[1].length_flattop2num_beta_osc() + upramp2.length_flattop2num_beta_osc() + downramp2.length_flattop2num_beta_osc(), 10.0, rtol=1e-3, atol=0.0) # Calculate the total phase adcance. + # Test plotting linac.stages[-1].plot_evolution(bunch='beam') linac.plot_survey() @@ -333,7 +382,7 @@ def test_ramped_linac(): # driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter) # main_source = setup_basic_main_source(plasma_density, ramp_beta_mag) -# stage = setup_StageReducedModels(plasma_density, driver_source, main_source, ramp_beta_mag, enable_tr_instability, enable_radiation_reaction, enable_ion_motion, use_ramps, drive_beam_update_period) +# stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, ramp_beta_mag=ramp_beta_mag, length_flattop=length_flattop, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps, drive_beam_update_period) # interstage = setup_InterstageImpactX(stage) # linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False) @@ -367,7 +416,7 @@ def test_ramped_linac(): # driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter) # main_source = setup_basic_main_source(plasma_density, ramp_beta_mag) -# stage = setup_StageReducedModels(plasma_density, driver_source, main_source, ramp_beta_mag, enable_tr_instability, enable_radiation_reaction, enable_ion_motion, use_ramps, drive_beam_update_period) +# stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, ramp_beta_mag=ramp_beta_mag, length_flattop=length_flattop, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps, drive_beam_update_period) # interstage = setup_InterstageImpactX(stage) # linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False) @@ -405,7 +454,7 @@ def test_angular_jitter_linac(): driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter) main_source = setup_basic_main_source(plasma_density, ramp_beta_mag) - stage = setup_StageReducedModels(plasma_density, driver_source, main_source, ramp_beta_mag, enable_tr_instability, enable_radiation_reaction, enable_ion_motion, use_ramps, drive_beam_update_period) + stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, ramp_beta_mag=ramp_beta_mag, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps) interstage = setup_InterstageImpactX(stage) linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False) @@ -439,11 +488,10 @@ def test_angular_jitter_ramped_linac(): enable_radiation_reaction = False enable_ion_motion = False use_ramps = True - drive_beam_update_period = 0 driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter) main_source = setup_basic_main_source(plasma_density, ramp_beta_mag) - stage = setup_StageReducedModels(plasma_density, driver_source, main_source, ramp_beta_mag, enable_tr_instability, enable_radiation_reaction, enable_ion_motion, use_ramps, drive_beam_update_period) + stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, ramp_beta_mag=ramp_beta_mag, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps) interstage = setup_InterstageImpactX(stage) linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False) @@ -477,12 +525,11 @@ def test_trInstability_linac(): enable_radiation_reaction = True enable_ion_motion = False use_ramps = False - drive_beam_update_period = 0 driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter) main_source = setup_basic_main_source(plasma_density, ramp_beta_mag) main_source.energy = 3.0e9 # [eV], HALHF v2 start energy - stage = setup_StageReducedModels(plasma_density, driver_source, main_source, ramp_beta_mag, enable_tr_instability, enable_radiation_reaction, enable_ion_motion, use_ramps, drive_beam_update_period) + stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, ramp_beta_mag=ramp_beta_mag, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps) interstage = setup_InterstageImpactX(stage) linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False) @@ -516,7 +563,7 @@ def test_jitter_trInstability_ramped_linac(): driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter) main_source = setup_basic_main_source(plasma_density, ramp_beta_mag, energy=3.0e9) - stage = setup_StageReducedModels(plasma_density, driver_source, main_source, ramp_beta_mag, enable_tr_instability, enable_radiation_reaction, enable_ion_motion, use_ramps, drive_beam_update_period) + stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, ramp_beta_mag=ramp_beta_mag, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps) interstage = setup_InterstageImpactX(stage) linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False) @@ -614,7 +661,7 @@ def test_jitter_trInstability_ramped_linac(): # driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter) # main_source = setup_basic_main_source(plasma_density, ramp_beta_mag) -# stage = setup_StageReducedModels(plasma_density, driver_source, main_source, ramp_beta_mag, enable_tr_instability, enable_radiation_reaction, enable_ion_motion, use_ramps, drive_beam_update_period) +# stage = setup_StageReducedModels(plasma_density, driver_source, ramp_beta_mag, enable_tr_instability, enable_radiation_reaction, enable_ion_motion, use_ramps, drive_beam_update_period) # interstage = setup_InterstageImpactX(stage) # linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False) @@ -644,11 +691,10 @@ def test_ionMotion_linac(): enable_radiation_reaction = True enable_ion_motion = True use_ramps = False - drive_beam_update_period = 0 driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter) main_source = setup_basic_main_source(plasma_density, ramp_beta_mag, energy=361.8e9) - stage = setup_StageReducedModels(plasma_density, driver_source, main_source, ramp_beta_mag, enable_tr_instability, enable_radiation_reaction, enable_ion_motion, use_ramps, drive_beam_update_period) + stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, ramp_beta_mag=ramp_beta_mag, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps) interstage = setup_InterstageImpactX(stage) linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False) @@ -679,11 +725,10 @@ def test_jitter_trInstability_ionMotion_linac(): enable_radiation_reaction = True enable_ion_motion = True use_ramps = False - drive_beam_update_period = 0 driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter) main_source = setup_basic_main_source(plasma_density, ramp_beta_mag, energy=361.8e9) - stage = setup_StageReducedModels(plasma_density, driver_source, main_source, ramp_beta_mag, enable_tr_instability, enable_radiation_reaction, enable_ion_motion, use_ramps, drive_beam_update_period) + stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, ramp_beta_mag=ramp_beta_mag, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps) interstage = setup_InterstageImpactX(stage) linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False) @@ -718,12 +763,11 @@ def test_jitter_trInstability_ionMotion_ramped_linac(): enable_radiation_reaction = True enable_ion_motion = True use_ramps = True - drive_beam_update_period = 0 driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter) assert driver_source.align_beam_axis is False main_source = setup_basic_main_source(plasma_density, ramp_beta_mag, energy=81.0e9) # Choosing an energy that gives a sensible number of time steps. - stage = setup_StageReducedModels(plasma_density, driver_source, main_source, ramp_beta_mag, enable_tr_instability, enable_radiation_reaction, enable_ion_motion, use_ramps, drive_beam_update_period, save_final_step=True) + stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, ramp_beta_mag=ramp_beta_mag, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps, save_final_step=True) interstage = setup_InterstageImpactX(stage) assert stage.driver_source.align_beam_axis is True @@ -824,7 +868,7 @@ def test_jitter_trInstability_ionMotion_ramped_linac(): # driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter) # main_source = setup_basic_main_source(plasma_density, ramp_beta_mag, energy=3.0e9) -# stage = setup_StageReducedModels(plasma_density, driver_source, main_source, ramp_beta_mag, enable_tr_instability, enable_radiation_reaction, enable_ion_motion, use_ramps, drive_beam_update_period) +# stage = setup_StageReducedModels(plasma_density, driver_source, ramp_beta_mag, enable_tr_instability, enable_radiation_reaction, enable_ion_motion, use_ramps, drive_beam_update_period) # interstage = setup_InterstageImpactX(stage) # linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False)