From dd147350f21b9cb928d5d6795c2ed175ba1531f3 Mon Sep 17 00:00:00 2001 From: Allessyer Date: Fri, 11 Feb 2022 18:45:52 +0100 Subject: [PATCH 1/4] Add rectangular trapezoidal shape of electron beam --- aptools/utilities/bunch_generation.py | 80 ++++++++++++++++++++++++++- 1 file changed, 79 insertions(+), 1 deletion(-) diff --git a/aptools/utilities/bunch_generation.py b/aptools/utilities/bunch_generation.py index b78a12c..8645c4b 100644 --- a/aptools/utilities/bunch_generation.py +++ b/aptools/utilities/bunch_generation.py @@ -11,7 +11,7 @@ def generate_gaussian_bunch_from_twiss( - a_x, a_y, b_x, b_y, en_x, en_y, ene, ene_sp, s_t, q_tot, n_part, x_c=0, + a_x, a_y, b_x, b_y, en_x, en_y, ene, ene_sp, s_t, q_tot, n_part, head_current, x_c=0, y_c=0, z_c=0, lon_profile='gauss', min_len_scale_noise=None, sigma_trunc_lon=None, smooth_sigma=None, smooth_trunc=None, save_to_file=False, save_to_code='astra', @@ -56,6 +56,10 @@ def generate_gaussian_bunch_from_twiss( n_part: int Total number of particles in the bunch. + + head_current: float + takes a value between 0 and 1. + if it is 1 all the current is in the head, if it 0 all the current is in the tail, and if it is 0.5 both head and tail have the same current. x_c: float Central bunch position in the x-plane in units of m. @@ -142,6 +146,10 @@ def generate_gaussian_bunch_from_twiss( elif lon_profile == 'flattop_smoothed': z = _create_flattop_longitudinal_profile_with_smoothing( z_c, s_z, n_part, min_len_scale_noise, smooth_sigma, smooth_trunc) + elif lon_profile == 'rectan_trapezoidal': + z = _create_rectan_trapezoidal_longitudinal_profile(z_c, s_z, n_part,head_current, + min_len_scale_noise) + # Define again n_part in case it changed when crealing long. profile n_part = len(z) pz = np.random.normal(ene, ene_sp_abs, n_part) @@ -308,6 +316,76 @@ def _create_flattop_longitudinal_profile_with_smoothing(z_c, length, n_part, z = z - length/2 + z_c return z +def _create_rectan_trapezoidal_longitudinal_profile(z_c, length, n_part, head_current, + min_len_scale_noise): + """ Creates a rectangular trapezoidal longitudinal profile """ + + # Make sure number of particles is an integer + n_part = int(n_part) + if min_len_scale_noise is None: + z = rectan_trapezoid(n_part,a=z_c-length/2,b=z_c+length/2,h1=head_current) + else: + raise NotImplementedError('Noise reduction not implemented for `trapezoidal` profile.') + return z + + +def rectan_trapezoid(n_part,a=0.,b=1.,h1=0.2): + + """ + Creates a longitudinal rectangular trapezoid particle bunch with the specified + parameters. + Parameters + ---------- + n_part : float + Number of particles in the beam + + __ + __/ | + __/ | + / | + | | + | | + |__________| + a b + + a : float + start position of the trapezoid + b : float + end position of the trapezoid + h1 : float + % from the maximum height + h2 : float + second height + plot : bool + If True, then plot histogram of the distribution, + + ---------- + Returns + + x : array with size [n_part] + distribution of random variables with rectangular trapezoidal shape + + """ + + n_part = int(n_part) + y = np.random.uniform(0,1,n_part) + + h1 = h1 * 2/(b-a) + + # Find h2 by condition of S = 1 (whole area is equal to 1, because it is probability) + h2 = 2/(b-a) - h1 + + x = np.zeros(len(y)) + + for i in range(len(y)): + if y[i] >= 0 and y[i] <= 1: + n = h1 + m = (h2-h1)/(2*(b-a)) + D = n**2 + 4 * m * y[i] + x[i] = (-n+np.sqrt(D))/(2*m) + a + + return x + def _create_smooth_z_array(part_per_slice, slice_edges): """ Creates the z array of the distribution when forced to be smooth """ From 2626ffa3543b70bbacefc4927bb02de48a1f850f Mon Sep 17 00:00:00 2001 From: Allessyer Date: Fri, 11 Feb 2022 19:07:23 +0100 Subject: [PATCH 2/4] Add rectangular trapezoidal shape of electron beam --- aptools/utilities/bunch_generation.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/aptools/utilities/bunch_generation.py b/aptools/utilities/bunch_generation.py index 8645c4b..09bd806 100644 --- a/aptools/utilities/bunch_generation.py +++ b/aptools/utilities/bunch_generation.py @@ -377,11 +377,11 @@ def rectan_trapezoid(n_part,a=0.,b=1.,h1=0.2): x = np.zeros(len(y)) - for i in range(len(y)): - if y[i] >= 0 and y[i] <= 1: + for i, y_i in enumerate(y): + if y_i >= 0 and y_i <= 1: n = h1 m = (h2-h1)/(2*(b-a)) - D = n**2 + 4 * m * y[i] + D = n**2 + 4 * m * y_i x[i] = (-n+np.sqrt(D))/(2*m) + a return x From c14baa796aa57b1e71611ccf37925697ffac824f Mon Sep 17 00:00:00 2001 From: Allessyer Date: Wed, 2 Mar 2022 04:39:22 +0100 Subject: [PATCH 3/4] Add a rectangular trapezoidal smoothed longitudinal profile --- aptools/utilities/bunch_generation.py | 189 +++++++++++++++++++++++--- 1 file changed, 173 insertions(+), 16 deletions(-) diff --git a/aptools/utilities/bunch_generation.py b/aptools/utilities/bunch_generation.py index 09bd806..6a780d8 100644 --- a/aptools/utilities/bunch_generation.py +++ b/aptools/utilities/bunch_generation.py @@ -149,6 +149,10 @@ def generate_gaussian_bunch_from_twiss( elif lon_profile == 'rectan_trapezoidal': z = _create_rectan_trapezoidal_longitudinal_profile(z_c, s_z, n_part,head_current, min_len_scale_noise) + elif lon_profile == 'rectan_trapezoidal_smoothed': + z = _create_rectan_trapezoidal_longitudinal_profile_smoothed(z_c, s_z, n_part, + head_current, + min_len_scale_noise) # Define again n_part in case it changed when crealing long. profile n_part = len(z) @@ -323,14 +327,43 @@ def _create_rectan_trapezoidal_longitudinal_profile(z_c, length, n_part, head_cu # Make sure number of particles is an integer n_part = int(n_part) if min_len_scale_noise is None: - z = rectan_trapezoid(n_part,a=z_c-length/2,b=z_c+length/2,h1=head_current) + a = z_c-length/2 + b = z_c+length/2 + + if head_current == 0.5: + head_current += 1e-12 + + hmax = 2/(b-a) + h1 = head_current * hmax + + h2 = 2/(b-a) - h1 + if h2 < 0: + raise ValueError("h2 = 2/(b-a) - h1 < 0!") + + z = rectan_trapezoid(n_part,a=a,b=b,h1=h1,h2=h2) else: raise NotImplementedError('Noise reduction not implemented for `trapezoidal` profile.') return z -def rectan_trapezoid(n_part,a=0.,b=1.,h1=0.2): +def _create_rectan_trapezoidal_longitudinal_profile_smoothed(z_c, length, n_part, head_current, + min_len_scale_noise, + left_weight=0.05,right_weight=0.05): + + """ Creates a rectangular trapezoidal smoothed longitudinal profile """ + + # Make sure number of particles is an integer + n_part = int(n_part) + if min_len_scale_noise is None: + z = rectan_trapezoid_smoothed(n_part=n_part,a=z_c-length/2,b=z_c+length/2, + h1=head_current,left_weight=left_weight, + right_weight=right_weight) + else: + raise NotImplementedError('Noise reduction not implemented for `trapezoidal` profile.') + return z + +def rectan_trapezoid(n_part,a,b,h1,h2,plot=False): """ Creates a longitudinal rectangular trapezoid particle bunch with the specified parameters. @@ -344,7 +377,7 @@ def rectan_trapezoid(n_part,a=0.,b=1.,h1=0.2): __/ | / | | | - | | + |h1 |h2 |__________| a b @@ -353,12 +386,11 @@ def rectan_trapezoid(n_part,a=0.,b=1.,h1=0.2): b : float end position of the trapezoid h1 : float - % from the maximum height + first height. h1 can only have a value between 0 and maximum height: 2/(b-a) h2 : float second height - plot : bool - If True, then plot histogram of the distribution, - + plot: bool + if True, then plot PDF ---------- Returns @@ -366,24 +398,149 @@ def rectan_trapezoid(n_part,a=0.,b=1.,h1=0.2): distribution of random variables with rectangular trapezoidal shape """ - n_part = int(n_part) y = np.random.uniform(0,1,n_part) - h1 = h1 * 2/(b-a) - - # Find h2 by condition of S = 1 (whole area is equal to 1, because it is probability) - h2 = 2/(b-a) - h1 - x = np.zeros(len(y)) - for i, y_i in enumerate(y): - if y_i >= 0 and y_i <= 1: + for i in range(len(y)): + if y[i] >= 0 and y[i] <= 1: n = h1 m = (h2-h1)/(2*(b-a)) - D = n**2 + 4 * m * y_i + D = n**2 + 4 * m * y[i] x[i] = (-n+np.sqrt(D))/(2*m) + a + if plot: + plt.hist(x,density=True,histtype='stepfilled',bins=50, alpha=0.2); + plt.grid() + # plt.ylim(0.,1.); + + return x + + +def half_gaussian(n_part, mu, peak, part): + """ + Creates half gaussian distribution + + + /| peak|\ + left -----> / | | \ <---- right gaussian distribution + gaussian / | | \ + distribution / | | \ + mu + + n_part : float + Number of particles in the distribution + mu : float + Mean of the distribution, if it would be whole gaussian distribution + peak : float + Height of the distribution + part : string + it can be only "left" or "right". + indicates which part of gaussian distribution you want to create. + ---------- + Returns + + x : array with size [n_part] + distribution of random variables with half gaussian shape + + """ + + if peak == 0: + peak += 1e-12 + + # Since we use all the particles and translate them either to the left or to the right, + # if the number doubles on the left/right, + # then the height is halved (or sigma is doubled) + + sigma = 1 / (peak / 2 * np.sqrt(2 * np.pi)) + + x = np.random.normal(loc=mu, scale=sigma, size=n_part) + if part == 'right': + idx = x < mu + elif part == 'left': + idx = x > mu + else: + raise ValueError('Wrong part') + x[idx] = 2 * mu - x[idx] + return x + +def rectan_trapezoid_smoothed(n_part,a,b,h1,left_weight,right_weight,plot=False): + """ + Creates a longitudinal rectangular trapezoid particle bunch with smoothed sides. + Parameters + ---------- + n_part : float + Number of particles in the beam + + ______ + __/ |\ + ___/ | \ + gaussian /| | \ + distribution / | trapezoial | \ gaussian distribution + 0.05% n_part--> / |distribution | \ <--0.05% n_part + /__ |__0.9n_part__|_____\ + a b + + a : float + start position of the trapezoid + b : float + end position of the trapezoid + h1 : float + first height + % of the maximum height + h2 : float + second height + + left_weight : float + % from the whole area which goes to the left half-gaussian distribution + + right_weight : float + % from the whole area which goes to the right half-gaussian distribution + + plot: bool + if True, then plot PDF + ---------- + Returns + + x : array with size [n_part] + distribution of random variables with rectangular trapezoidal smoothed shape + + """ + + trapez_weight = 1 - (left_weight + right_weight) + n_part = int(n_part) + + if h1 == 0.5: + h1 += 1e-12 + elif h1 < 0.01 or h1 > 0.99: + raise ValueError('Please, put head_current in range [0.01,0.99]') + + hmax = 2/(b-a) * trapez_weight + h1 = h1 * hmax + + h2 = hmax - h1 + if h2 < 0: + raise ValueError("h2 = 2/(b-a) - h1 < 0!") + + p = np.random.rand(n_part) + x = np.zeros_like(p) + + idx_left = p < left_weight + idx_right = p > (1 - right_weight) + idx = np.logical_not(np.logical_or(idx_left, idx_right)) + + if trapez_weight != 0: + x[idx] = rectan_trapezoid(np.count_nonzero(idx), a, b, h1 / trapez_weight, h2 / trapez_weight,plot=False) + if left_weight != 0: + x[idx_left] = half_gaussian(np.count_nonzero(idx_left), mu=a, peak=h1 / left_weight, part='left') + if right_weight != 0: + x[idx_right] = half_gaussian(np.count_nonzero(idx_right), mu=b, peak=h2 / right_weight, part='right') + + if plot: + plt.hist(x,histtype='stepfilled', bins=100, alpha=1., density=True) + plt.grid() + return x From 75a343f1464a49ef8faacc1e4423086c35942291 Mon Sep 17 00:00:00 2001 From: Allessyer Date: Fri, 4 Mar 2022 15:20:11 +0100 Subject: [PATCH 4/4] Update rectangular trapezoidal smoothed shape --- aptools/utilities/bunch_generation.py | 65 +++++++++++++-------------- 1 file changed, 32 insertions(+), 33 deletions(-) diff --git a/aptools/utilities/bunch_generation.py b/aptools/utilities/bunch_generation.py index 6a780d8..de198fc 100644 --- a/aptools/utilities/bunch_generation.py +++ b/aptools/utilities/bunch_generation.py @@ -11,8 +11,8 @@ def generate_gaussian_bunch_from_twiss( - a_x, a_y, b_x, b_y, en_x, en_y, ene, ene_sp, s_t, q_tot, n_part, head_current, x_c=0, - y_c=0, z_c=0, lon_profile='gauss', min_len_scale_noise=None, + a_x, a_y, b_x, b_y, en_x, en_y, ene, ene_sp, s_t, q_tot, n_part, head_current, + x_c=0, y_c=0, z_c=0, lon_profile='gauss', min_len_scale_noise=None, sigma_trunc_lon=None, smooth_sigma=None, smooth_trunc=None, save_to_file=False, save_to_code='astra', save_to_path=None, file_name=None, perform_checks=False): @@ -59,7 +59,8 @@ def generate_gaussian_bunch_from_twiss( head_current: float takes a value between 0 and 1. - if it is 1 all the current is in the head, if it 0 all the current is in the tail, and if it is 0.5 both head and tail have the same current. + if it is 1 all the current is in the head, if it 0 all the current is + in the tail, and if it is 0.5 both head and tail have the same current. x_c: float Central bunch position in the x-plane in units of m. @@ -147,12 +148,15 @@ def generate_gaussian_bunch_from_twiss( z = _create_flattop_longitudinal_profile_with_smoothing( z_c, s_z, n_part, min_len_scale_noise, smooth_sigma, smooth_trunc) elif lon_profile == 'rectan_trapezoidal': - z = _create_rectan_trapezoidal_longitudinal_profile(z_c, s_z, n_part,head_current, - min_len_scale_noise) + z = _create_rectan_trapezoidal_longitudinal_profile(z_c, s_z, n_part, + head_current, + min_len_scale_noise) elif lon_profile == 'rectan_trapezoidal_smoothed': - z = _create_rectan_trapezoidal_longitudinal_profile_smoothed(z_c, s_z, n_part, - head_current, - min_len_scale_noise) + z = _create_rectan_trapezoidal_longitudinal_profile_smoothed(z_c, + s_z, + n_part, + head_current, + min_len_scale_noise) # Define again n_part in case it changed when crealing long. profile n_part = len(z) @@ -205,7 +209,6 @@ def generate_from_file_modifying_twiss( read_kwargs: dict Dictionary containing optional parameters for the read_beam function. - save_to_file: bool Whether to save the generated distribution to a file. @@ -320,8 +323,10 @@ def _create_flattop_longitudinal_profile_with_smoothing(z_c, length, n_part, z = z - length/2 + z_c return z -def _create_rectan_trapezoidal_longitudinal_profile(z_c, length, n_part, head_current, - min_len_scale_noise): + +def _create_rectan_trapezoidal_longitudinal_profile(z_c, length, n_part, + head_current, + min_len_scale_noise): """ Creates a rectangular trapezoidal longitudinal profile """ # Make sure number of particles is an integer @@ -346,9 +351,11 @@ def _create_rectan_trapezoidal_longitudinal_profile(z_c, length, n_part, head_cu return z -def _create_rectan_trapezoidal_longitudinal_profile_smoothed(z_c, length, n_part, head_current, +def _create_rectan_trapezoidal_longitudinal_profile_smoothed(z_c, length, + n_part, head_current, min_len_scale_noise, - left_weight=0.05,right_weight=0.05): + left_weight=0.05, + right_weight=0.05): """ Creates a rectangular trapezoidal smoothed longitudinal profile """ @@ -363,7 +370,7 @@ def _create_rectan_trapezoidal_longitudinal_profile_smoothed(z_c, length, n_part return z -def rectan_trapezoid(n_part,a,b,h1,h2,plot=False): +def rectan_trapezoid(n_part, a, b, h1, h2): """ Creates a longitudinal rectangular trapezoid particle bunch with the specified parameters. @@ -389,8 +396,7 @@ def rectan_trapezoid(n_part,a,b,h1,h2,plot=False): first height. h1 can only have a value between 0 and maximum height: 2/(b-a) h2 : float second height - plot: bool - if True, then plot PDF + ---------- Returns @@ -409,12 +415,6 @@ def rectan_trapezoid(n_part,a,b,h1,h2,plot=False): m = (h2-h1)/(2*(b-a)) D = n**2 + 4 * m * y[i] x[i] = (-n+np.sqrt(D))/(2*m) + a - - if plot: - plt.hist(x,density=True,histtype='stepfilled',bins=50, alpha=0.2); - plt.grid() - # plt.ylim(0.,1.); - return x @@ -465,7 +465,8 @@ def half_gaussian(n_part, mu, peak, part): x[idx] = 2 * mu - x[idx] return x -def rectan_trapezoid_smoothed(n_part,a,b,h1,left_weight,right_weight,plot=False): + +def rectan_trapezoid_smoothed(n_part, a, b, h1, left_weight, right_weight): """ Creates a longitudinal rectangular trapezoid particle bunch with smoothed sides. Parameters @@ -504,8 +505,7 @@ def rectan_trapezoid_smoothed(n_part,a,b,h1,left_weight,right_weight,plot=False) Returns x : array with size [n_part] - distribution of random variables with rectangular trapezoidal smoothed shape - + distribution of random variables with rectangular trapezoidal smoothed shape """ trapez_weight = 1 - (left_weight + right_weight) @@ -531,16 +531,15 @@ def rectan_trapezoid_smoothed(n_part,a,b,h1,left_weight,right_weight,plot=False) idx = np.logical_not(np.logical_or(idx_left, idx_right)) if trapez_weight != 0: - x[idx] = rectan_trapezoid(np.count_nonzero(idx), a, b, h1 / trapez_weight, h2 / trapez_weight,plot=False) + x[idx] = rectan_trapezoid(np.count_nonzero(idx), a, b, + h1 / trapez_weight, + h2 / trapez_weight,plot=False) if left_weight != 0: - x[idx_left] = half_gaussian(np.count_nonzero(idx_left), mu=a, peak=h1 / left_weight, part='left') + x[idx_left] = half_gaussian(np.count_nonzero(idx_left), mu=a, + peak=h1 / left_weight, part='left') if right_weight != 0: - x[idx_right] = half_gaussian(np.count_nonzero(idx_right), mu=b, peak=h2 / right_weight, part='right') - - if plot: - plt.hist(x,histtype='stepfilled', bins=100, alpha=1., density=True) - plt.grid() - + x[idx_right] = half_gaussian(np.count_nonzero(idx_right), mu=b, + peak=h2 / right_weight, part='right') return x