From 561a3c40ae336377f0c3794393f6f5ea17aeec28 Mon Sep 17 00:00:00 2001 From: elwaagaa Date: Mon, 5 Dec 2022 16:32:37 +0100 Subject: [PATCH 1/8] Added ion scaling tracking tests comparing maximum SC tune shifts w.r.t protons. --- tests/test_ion_scaling.py | 182 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 182 insertions(+) create mode 100644 tests/test_ion_scaling.py diff --git a/tests/test_ion_scaling.py b/tests/test_ion_scaling.py new file mode 100644 index 0000000..fb1a691 --- /dev/null +++ b/tests/test_ion_scaling.py @@ -0,0 +1,182 @@ +# copyright ############################### # +# This file is part of the Xpart Package. # +# Copyright (c) CERN, 2021. # +# ######################################### # + +import numpy as np +import json +import NAFFlib + +import xobjects as xo +import xpart as xp +import xtrack as xt +import xfields as xf + +test_data_folder = xt._pkg_root.joinpath('../test_data').absolute() + +def test_ion_scaling_tracking(): + """ + Test if the Q^2/A = 1 scaling holds versus a proton reference case + for some linearly increasing values of Q^2 and A + """ + for context in xo.context.get_test_contexts(): + print(f"Test {context.__class__}") + ########## Test settings ################## + bunch_intensity = 1e11/3 # Need short bunch to avoid bucket non-linearity + sigma_z = 22.5e-2/3 + nemitt_x = 2.5e-6 + nemitt_y = 2.5e-6 + num_particles = 1000 + num_turns = 30 + + num_spacecharge_interactions = 540 + tol_spacecharge_position = 1e-2 + + ########## Load the SPS proton sequence and ref particle ##################### + fname_line = test_data_folder.joinpath('sps_w_spacecharge/line_no_spacecharge_and_particle.json') + with open(fname_line, 'r') as fid: + input_data = json.load(fid) + + # Load line for tracking and reference particle + line = xt.Line.from_dict(input_data['line']) + particle_ref = xp.Particles.from_dict(input_data['particle']) + ctx2arr = context.nparray_from_context_array # Copies an array to the device to a numpy array. + + # Define longitudinal profile + lprofile = xf.LongitudinalProfileQGaussian( + number_of_particles=bunch_intensity, + sigma_z=sigma_z, + z0=0., + q_parameter=1.) + + ########### PROTON REFERENCE CASE ############# + # Install space charge + xf.install_spacecharge_frozen(line=line, + particle_ref=particle_ref, + longitudinal_profile=lprofile, + nemitt_x=nemitt_x, nemitt_y=nemitt_y, + sigma_z=sigma_z, + num_spacecharge_interactions=num_spacecharge_interactions, + tol_spacecharge_position=tol_spacecharge_position) + + + # Define tracker and particle beam + tracker = xt.Tracker(_context=context, + line=line) + tracker_sc_off = tracker.filter_elements(exclude_types_starting_with='SpaceCh') + + particles0 = xp.generate_matched_gaussian_bunch(_context=context, + num_particles=num_particles, total_intensity_particles=bunch_intensity, + nemitt_x=nemitt_x, nemitt_y=nemitt_y, sigma_z=sigma_z, + particle_ref=particle_ref, tracker=tracker_sc_off) + + # Set up for the tracking + tracker.optimize_for_tracking() + x_tbt = np.zeros((num_particles, num_turns), dtype=np.float64) + y_tbt = np.zeros((num_particles, num_turns), dtype=np.float64) + Qx0 = np.zeros(num_particles) + Qy0 = np.zeros(num_particles) + + # Perform tracking + for ii in range(num_turns): + print(f'Turn: {ii}\n', end='\r', flush=True) + x_tbt[:, ii] = ctx2arr(particles0.x[:num_particles]).copy() + y_tbt[:, ii] = ctx2arr(particles0.y[:num_particles]).copy() + tracker.track(particles0) + + # Find maximum tune shift + for i_part in range(num_particles): + Qx0[i_part] = NAFFlib.get_tune(x_tbt[i_part, :]) + Qy0[i_part] = NAFFlib.get_tune(y_tbt[i_part, :]) + + Qx0_max = np.max(Qx0) + Qy0_max = np.max(Qy0) + + ################################################### + + ############ ION SCALING TRACKING ################# + proton_ref_momentum = 25.92e9 + As = np.array([4., 9., 25., 100., 144.]) # mass numbers + Qs = np.array([2., 3., 5., 10., 12.]) # charge states + Qx_max_array = [] + Qy_max_array = [] + + for jj in range(len(As)): + print("\nIon scaling nr {}\n".format(jj+1)) + x_tbt_ion = np.zeros((num_particles, num_turns), dtype=np.float64) + y_tbt_ion = np.zeros((num_particles, num_turns), dtype=np.float64) + Qx = np.zeros(num_particles) + Qy = np.zeros(num_particles) + + + # Modify reference particle according to new masses + mass_ion = As[jj]*931494102.42 # atomic mass unit-electron volt relationship + charge_ion = Qs[jj] + + line_ion = xt.Line.from_dict(input_data['line']) + + particle_sample = xp.Particles( + mass0 = mass_ion, + q0= charge_ion, + p0c = proton_ref_momentum*As[jj] + ) + + ########### INSTALL SPACE CHARGE FOR IONS ################# + xf.install_spacecharge_frozen(line=line_ion, + particle_ref=particle_sample, + longitudinal_profile=lprofile, + nemitt_x=nemitt_x, nemitt_y=nemitt_y, + sigma_z=sigma_z, + num_spacecharge_interactions=num_spacecharge_interactions, + tol_spacecharge_position=tol_spacecharge_position) + + + ############ DEFINE TRACKER AND PARTICLE BEAM ###################### + tracker_ion = xt.Tracker(_context=context, + line=line_ion) + tracker_sc_off_ion = tracker_ion.filter_elements(exclude_types_starting_with='SpaceCh') + + particles_ion = xp.generate_matched_gaussian_bunch(_context=context, + num_particles=num_particles, total_intensity_particles=bunch_intensity, + nemitt_x=nemitt_x, nemitt_y=nemitt_y, sigma_z=sigma_z, + particle_ref=particle_sample, tracker=tracker_sc_off_ion) + + + ########### TRACK THE PARTICLES ################ + tracker_ion.optimize_for_tracking() + + # Perform tracking + for ii in range(num_turns): + print(f'Turn: {ii}\n', end='\r', flush=True) + x_tbt_ion[:, ii] = ctx2arr(particles_ion.x[:num_particles]).copy() + y_tbt_ion[:, ii] = ctx2arr(particles_ion.y[:num_particles]).copy() + tracker_ion.track(particles_ion) + + # Find maximum tune shift + for i_part in range(num_particles): + Qx[i_part] = NAFFlib.get_tune(x_tbt_ion[i_part, :]) + Qy[i_part] = NAFFlib.get_tune(y_tbt_ion[i_part, :]) + + # Remove any possible outliers due to resonances + Qx = Qx[(Qx < 1.5*Qx0_max) & (Qx > 0.5*np.min(Qx0))] + Qy = Qy[(Qy < 1.5*Qy0_max) & (Qy > 0.5*np.min(Qy0))] + + # Check maximum tune shift + Qx_max = np.max(Qx) + Qy_max = np.max(Qy) + Qx_max_array.append(Qx_max) + Qy_max_array.append(Qy_max) + print("Ion type {}: dQx = {}, dQy = {}".format(jj+1, Qx_max, Qy_max)) + print("Proton: dQx = {}, dQy = {}".format(Qx0_max, Qy0_max)) + + assert np.isclose(Qx_max, Qx0_max, atol=1e-2) + + +#fname_line = '../../xtrack/test_data/sps_w_spacecharge/line_no_spacecharge_and_particle.json' +#context = xo.ContextCpu() + + + + + + From 6f6a808884312ddb9e6c14d86caf802ee2080a47 Mon Sep 17 00:00:00 2001 From: elwaagaa Date: Mon, 5 Dec 2022 16:33:49 +0100 Subject: [PATCH 2/8] Remove unnecessary spaces. --- tests/test_ion_scaling.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/tests/test_ion_scaling.py b/tests/test_ion_scaling.py index fb1a691..29c7df0 100644 --- a/tests/test_ion_scaling.py +++ b/tests/test_ion_scaling.py @@ -32,7 +32,7 @@ def test_ion_scaling_tracking(): num_spacecharge_interactions = 540 tol_spacecharge_position = 1e-2 - ########## Load the SPS proton sequence and ref particle ##################### + ########## Load the SPS proton sequence and ref particle ##################### fname_line = test_data_folder.joinpath('sps_w_spacecharge/line_no_spacecharge_and_particle.json') with open(fname_line, 'r') as fid: input_data = json.load(fid) @@ -172,10 +172,6 @@ def test_ion_scaling_tracking(): assert np.isclose(Qx_max, Qx0_max, atol=1e-2) -#fname_line = '../../xtrack/test_data/sps_w_spacecharge/line_no_spacecharge_and_particle.json' -#context = xo.ContextCpu() - - From 2178237b427aa51abd7dc4560b1f20b01440deb7 Mon Sep 17 00:00:00 2001 From: elwaagaa Date: Mon, 5 Dec 2022 17:01:25 +0100 Subject: [PATCH 3/8] Further removed extra spaces and indents. --- tests/test_ion_scaling.py | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/tests/test_ion_scaling.py b/tests/test_ion_scaling.py index 29c7df0..1ce6cc2 100644 --- a/tests/test_ion_scaling.py +++ b/tests/test_ion_scaling.py @@ -169,10 +169,4 @@ def test_ion_scaling_tracking(): print("Ion type {}: dQx = {}, dQy = {}".format(jj+1, Qx_max, Qy_max)) print("Proton: dQx = {}, dQy = {}".format(Qx0_max, Qy0_max)) - assert np.isclose(Qx_max, Qx0_max, atol=1e-2) - - - - - - + assert np.isclose(Qx_max, Qx0_max, atol=1e-2) \ No newline at end of file From 76a9904b6db188dbbe35286ce6d0e3b5118d41e6 Mon Sep 17 00:00:00 2001 From: elwaagaa Date: Tue, 23 May 2023 19:28:38 +0200 Subject: [PATCH 4/8] Function and test to generate longitudinal particle distribution. --- tests/test_build_particles_parabolic.py | 51 +++++++++++++ ...ate_parabolic_longitudinal_distribution.py | 74 +++++++++++++++++++ 2 files changed, 125 insertions(+) create mode 100644 tests/test_build_particles_parabolic.py create mode 100644 xpart/longitudinal/generate_parabolic_longitudinal_distribution.py diff --git a/tests/test_build_particles_parabolic.py b/tests/test_build_particles_parabolic.py new file mode 100644 index 0000000..6a97635 --- /dev/null +++ b/tests/test_build_particles_parabolic.py @@ -0,0 +1,51 @@ +# copyright ############################### # +# This file is part of the Xpart Package. # +# Copyright (c) CERN, 2021. # +# ######################################### # + +import json + +import numpy as np + +import xpart as xp +import xtrack as xt +import xobjects as xo + +from xpart.longitudinal.generate_parabolic_longitudinal_distribution import parabolic_longitudinal_distribution + +from xobjects.test_helpers import for_all_test_contexts + + +@for_all_test_contexts +def test_build_particles_parabolic(test_context): + for ctx_ref in [test_context, None]: + # Build a reference particle + p0 = xp.Particles(mass0=xp.PROTON_MASS_EV, q0=1, p0c=7e12, x=1, y=3, + delta=[10], _context=ctx_ref) + + + # Load machine model (from pymask) + filename = xt._pkg_root.parent.joinpath('test_data/lhc_no_bb/line_and_particle.json') + with open(filename, 'r') as fid: + input_data = json.load(fid) + line = xt.Line.from_dict(input_data['line']) + line.build_tracker(_context=test_context) + + # Built a set of three particles with different x coordinates + particles = parabolic_longitudinal_distribution(_context=test_context, + num_particles=100, + nemitt_x=3e-6, + nemitt_y=3e-6, + sigma_z=0.05, + particle_ref=p0, + total_intensity_particles=1e10, + line=line + ) + + dct = particles.to_dict() + assert np.all(dct['p0c'] == 7e12) + tw = line.twiss(particle_ref=p0) + + + + diff --git a/xpart/longitudinal/generate_parabolic_longitudinal_distribution.py b/xpart/longitudinal/generate_parabolic_longitudinal_distribution.py new file mode 100644 index 0000000..8f33703 --- /dev/null +++ b/xpart/longitudinal/generate_parabolic_longitudinal_distribution.py @@ -0,0 +1,74 @@ +from xpart.longitudinal import generate_longitudinal_coordinates +from xpart import build_particles +import numpy as np + +def parabolic_longitudinal_distribution(_context=None, + num_particles=None, + nemitt_x=None, + nemitt_y=None, + sigma_z=None, + particle_ref=None, + total_intensity_particles=None, + tracker=None, + line=None + ): + + """ + Function to generate a parabolic longitudinal distribution + """ + + if line is not None and tracker is not None: + raise ValueError( + 'line and tracker cannot be provided at the same time.') + + if tracker is not None: + _print('Warning! ' + "The argument tracker is deprecated. Please use line instead.") + line = tracker.line + + if line is not None: + assert line.tracker is not None, ("The line must have a tracker, " + "please call Line.build_tracker() first.") + + if num_particles is None: + raise ValueError( + 'Number of particles must be provided') + if sigma_z is None: + raise ValueError( + 'Bunch length sigma_z must be provided') + + if particle_ref is None: + raise ValueError( + 'Reference particle must be provided') + + # If emittances are not provided, set them to default value of one + if nemitt_x is None: + nemitt_x = 1.0 + + if nemitt_y is None: + nemitt_y = 1.0 + + # Generate longitudinal coordinates s + zeta, delta = generate_longitudinal_coordinates(line=line, distribution='parabolic', + num_particles=num_particles, + engine='single-rf-harmonic', sigma_z=sigma_z, + particle_ref=particle_ref) + + # Initiate normalized coordinates + x_norm = np.random.normal(size=num_particles) + px_norm = np.random.normal(size=num_particles) + y_norm = np.random.normal(size=num_particles) + py_norm = np.random.normal(size=num_particles) + + # If not provided, use number of particles as intensity + if total_intensity_particles is None: + total_intensity_particles = num_particles + + particles = build_particles(_context=None, particle_ref=particle_ref, + zeta=zeta, delta=delta, + x_norm=x_norm, px_norm=px_norm, + y_norm=y_norm, py_norm=py_norm, + nemitt_x=nemitt_x, nemitt_y=nemitt_y, + weight=total_intensity_particles/num_particles, line=line) + + return particles From 21d357feef9d22268338ddd3cc5f285c58024888 Mon Sep 17 00:00:00 2001 From: elwaagaa Date: Tue, 23 May 2023 19:35:54 +0200 Subject: [PATCH 5/8] removed outdated ion scaling from this branch. --- tests/test_ion_scaling.py | 172 -------------------------------------- 1 file changed, 172 deletions(-) delete mode 100644 tests/test_ion_scaling.py diff --git a/tests/test_ion_scaling.py b/tests/test_ion_scaling.py deleted file mode 100644 index 1ce6cc2..0000000 --- a/tests/test_ion_scaling.py +++ /dev/null @@ -1,172 +0,0 @@ -# copyright ############################### # -# This file is part of the Xpart Package. # -# Copyright (c) CERN, 2021. # -# ######################################### # - -import numpy as np -import json -import NAFFlib - -import xobjects as xo -import xpart as xp -import xtrack as xt -import xfields as xf - -test_data_folder = xt._pkg_root.joinpath('../test_data').absolute() - -def test_ion_scaling_tracking(): - """ - Test if the Q^2/A = 1 scaling holds versus a proton reference case - for some linearly increasing values of Q^2 and A - """ - for context in xo.context.get_test_contexts(): - print(f"Test {context.__class__}") - ########## Test settings ################## - bunch_intensity = 1e11/3 # Need short bunch to avoid bucket non-linearity - sigma_z = 22.5e-2/3 - nemitt_x = 2.5e-6 - nemitt_y = 2.5e-6 - num_particles = 1000 - num_turns = 30 - - num_spacecharge_interactions = 540 - tol_spacecharge_position = 1e-2 - - ########## Load the SPS proton sequence and ref particle ##################### - fname_line = test_data_folder.joinpath('sps_w_spacecharge/line_no_spacecharge_and_particle.json') - with open(fname_line, 'r') as fid: - input_data = json.load(fid) - - # Load line for tracking and reference particle - line = xt.Line.from_dict(input_data['line']) - particle_ref = xp.Particles.from_dict(input_data['particle']) - ctx2arr = context.nparray_from_context_array # Copies an array to the device to a numpy array. - - # Define longitudinal profile - lprofile = xf.LongitudinalProfileQGaussian( - number_of_particles=bunch_intensity, - sigma_z=sigma_z, - z0=0., - q_parameter=1.) - - ########### PROTON REFERENCE CASE ############# - # Install space charge - xf.install_spacecharge_frozen(line=line, - particle_ref=particle_ref, - longitudinal_profile=lprofile, - nemitt_x=nemitt_x, nemitt_y=nemitt_y, - sigma_z=sigma_z, - num_spacecharge_interactions=num_spacecharge_interactions, - tol_spacecharge_position=tol_spacecharge_position) - - - # Define tracker and particle beam - tracker = xt.Tracker(_context=context, - line=line) - tracker_sc_off = tracker.filter_elements(exclude_types_starting_with='SpaceCh') - - particles0 = xp.generate_matched_gaussian_bunch(_context=context, - num_particles=num_particles, total_intensity_particles=bunch_intensity, - nemitt_x=nemitt_x, nemitt_y=nemitt_y, sigma_z=sigma_z, - particle_ref=particle_ref, tracker=tracker_sc_off) - - # Set up for the tracking - tracker.optimize_for_tracking() - x_tbt = np.zeros((num_particles, num_turns), dtype=np.float64) - y_tbt = np.zeros((num_particles, num_turns), dtype=np.float64) - Qx0 = np.zeros(num_particles) - Qy0 = np.zeros(num_particles) - - # Perform tracking - for ii in range(num_turns): - print(f'Turn: {ii}\n', end='\r', flush=True) - x_tbt[:, ii] = ctx2arr(particles0.x[:num_particles]).copy() - y_tbt[:, ii] = ctx2arr(particles0.y[:num_particles]).copy() - tracker.track(particles0) - - # Find maximum tune shift - for i_part in range(num_particles): - Qx0[i_part] = NAFFlib.get_tune(x_tbt[i_part, :]) - Qy0[i_part] = NAFFlib.get_tune(y_tbt[i_part, :]) - - Qx0_max = np.max(Qx0) - Qy0_max = np.max(Qy0) - - ################################################### - - ############ ION SCALING TRACKING ################# - proton_ref_momentum = 25.92e9 - As = np.array([4., 9., 25., 100., 144.]) # mass numbers - Qs = np.array([2., 3., 5., 10., 12.]) # charge states - Qx_max_array = [] - Qy_max_array = [] - - for jj in range(len(As)): - print("\nIon scaling nr {}\n".format(jj+1)) - x_tbt_ion = np.zeros((num_particles, num_turns), dtype=np.float64) - y_tbt_ion = np.zeros((num_particles, num_turns), dtype=np.float64) - Qx = np.zeros(num_particles) - Qy = np.zeros(num_particles) - - - # Modify reference particle according to new masses - mass_ion = As[jj]*931494102.42 # atomic mass unit-electron volt relationship - charge_ion = Qs[jj] - - line_ion = xt.Line.from_dict(input_data['line']) - - particle_sample = xp.Particles( - mass0 = mass_ion, - q0= charge_ion, - p0c = proton_ref_momentum*As[jj] - ) - - ########### INSTALL SPACE CHARGE FOR IONS ################# - xf.install_spacecharge_frozen(line=line_ion, - particle_ref=particle_sample, - longitudinal_profile=lprofile, - nemitt_x=nemitt_x, nemitt_y=nemitt_y, - sigma_z=sigma_z, - num_spacecharge_interactions=num_spacecharge_interactions, - tol_spacecharge_position=tol_spacecharge_position) - - - ############ DEFINE TRACKER AND PARTICLE BEAM ###################### - tracker_ion = xt.Tracker(_context=context, - line=line_ion) - tracker_sc_off_ion = tracker_ion.filter_elements(exclude_types_starting_with='SpaceCh') - - particles_ion = xp.generate_matched_gaussian_bunch(_context=context, - num_particles=num_particles, total_intensity_particles=bunch_intensity, - nemitt_x=nemitt_x, nemitt_y=nemitt_y, sigma_z=sigma_z, - particle_ref=particle_sample, tracker=tracker_sc_off_ion) - - - ########### TRACK THE PARTICLES ################ - tracker_ion.optimize_for_tracking() - - # Perform tracking - for ii in range(num_turns): - print(f'Turn: {ii}\n', end='\r', flush=True) - x_tbt_ion[:, ii] = ctx2arr(particles_ion.x[:num_particles]).copy() - y_tbt_ion[:, ii] = ctx2arr(particles_ion.y[:num_particles]).copy() - tracker_ion.track(particles_ion) - - # Find maximum tune shift - for i_part in range(num_particles): - Qx[i_part] = NAFFlib.get_tune(x_tbt_ion[i_part, :]) - Qy[i_part] = NAFFlib.get_tune(y_tbt_ion[i_part, :]) - - # Remove any possible outliers due to resonances - Qx = Qx[(Qx < 1.5*Qx0_max) & (Qx > 0.5*np.min(Qx0))] - Qy = Qy[(Qy < 1.5*Qy0_max) & (Qy > 0.5*np.min(Qy0))] - - # Check maximum tune shift - Qx_max = np.max(Qx) - Qy_max = np.max(Qy) - Qx_max_array.append(Qx_max) - Qy_max_array.append(Qy_max) - print("Ion type {}: dQx = {}, dQy = {}".format(jj+1, Qx_max, Qy_max)) - print("Proton: dQx = {}, dQy = {}".format(Qx0_max, Qy0_max)) - - assert np.isclose(Qx_max, Qx0_max, atol=1e-2) \ No newline at end of file From 8799f93fe6741d6013d6555291e37c753ae8af4c Mon Sep 17 00:00:00 2001 From: elwaagaa Date: Tue, 30 May 2023 16:36:08 +0200 Subject: [PATCH 6/8] Updated test to check particle distribution against SingleRFHarmonicMatcher. --- tests/test_build_particles_parabolic.py | 39 +++++++++++++------ xpart/__init__.py | 1 + xpart/longitudinal/__init__.py | 1 + ...ate_parabolic_longitudinal_distribution.py | 12 ++++-- 4 files changed, 37 insertions(+), 16 deletions(-) diff --git a/tests/test_build_particles_parabolic.py b/tests/test_build_particles_parabolic.py index 6a97635..064d43c 100644 --- a/tests/test_build_particles_parabolic.py +++ b/tests/test_build_particles_parabolic.py @@ -11,7 +11,7 @@ import xtrack as xt import xobjects as xo -from xpart.longitudinal.generate_parabolic_longitudinal_distribution import parabolic_longitudinal_distribution +from xpart.longitudinal import parabolic_longitudinal_distribution from xobjects.test_helpers import for_all_test_contexts @@ -23,6 +23,9 @@ def test_build_particles_parabolic(test_context): p0 = xp.Particles(mass0=xp.PROTON_MASS_EV, q0=1, p0c=7e12, x=1, y=3, delta=[10], _context=ctx_ref) + # Parameters for the test + num_part = 1000000 + parabolic_parameter = 0.05 # Load machine model (from pymask) filename = xt._pkg_root.parent.joinpath('test_data/lhc_no_bb/line_and_particle.json') @@ -30,22 +33,34 @@ def test_build_particles_parabolic(test_context): input_data = json.load(fid) line = xt.Line.from_dict(input_data['line']) line.build_tracker(_context=test_context) - - # Built a set of three particles with different x coordinates - particles = parabolic_longitudinal_distribution(_context=test_context, - num_particles=100, - nemitt_x=3e-6, - nemitt_y=3e-6, - sigma_z=0.05, - particle_ref=p0, - total_intensity_particles=1e10, - line=line - ) + + # Built a set of three particles with different x coordinates + particles, matcher = parabolic_longitudinal_distribution( + num_particles=num_part, + nemitt_x=3e-6, + nemitt_y=3e-6, + sigma_z=parabolic_parameter, + particle_ref=p0, + total_intensity_particles=1e10, + line=line, + return_matcher=True + ) dct = particles.to_dict() assert np.all(dct['p0c'] == 7e12) tw = line.twiss(particle_ref=p0) + # Test if longitudinal coordinates match with Single + # Generate distribution from RF matcher + tau = particles.zeta / p0.beta0[0] + tau_distr_y = matcher.tau_distr_y + tau_distr_x = matcher.tau_distr_x + dx = tau_distr_x[1] - tau_distr_x[0] + hist, edges = np.histogram(tau, + range=(tau_distr_x[0]-dx/2., tau_distr_x[-1]+dx/2.), + bins=len(tau_distr_x)) + hist = hist / sum(hist) * sum(tau_distr_y) + assert np.all(np.isclose(hist, tau_distr_y, atol=5.e-2, rtol=1.e-2)) diff --git a/xpart/__init__.py b/xpart/__init__.py index 03bb0b8..dc6f1f1 100644 --- a/xpart/__init__.py +++ b/xpart/__init__.py @@ -16,6 +16,7 @@ from .transverse_generators import generate_2D_gaussian from .longitudinal import generate_longitudinal_coordinates +from .longitudinal import parabolic_longitudinal_distribution from .constants import PROTON_MASS_EV, ELECTRON_MASS_EV diff --git a/xpart/longitudinal/__init__.py b/xpart/longitudinal/__init__.py index 086f0b4..0a4fcbc 100644 --- a/xpart/longitudinal/__init__.py +++ b/xpart/longitudinal/__init__.py @@ -5,3 +5,4 @@ from .generate_longitudinal import generate_longitudinal_coordinates from .single_rf_harmonic_matcher import SingleRFHarmonicMatcher +from .generate_parabolic_longitudinal_distribution import parabolic_longitudinal_distribution diff --git a/xpart/longitudinal/generate_parabolic_longitudinal_distribution.py b/xpart/longitudinal/generate_parabolic_longitudinal_distribution.py index 8f33703..d84402b 100644 --- a/xpart/longitudinal/generate_parabolic_longitudinal_distribution.py +++ b/xpart/longitudinal/generate_parabolic_longitudinal_distribution.py @@ -10,7 +10,8 @@ def parabolic_longitudinal_distribution(_context=None, particle_ref=None, total_intensity_particles=None, tracker=None, - line=None + line=None, + return_matcher=False ): """ @@ -49,10 +50,10 @@ def parabolic_longitudinal_distribution(_context=None, nemitt_y = 1.0 # Generate longitudinal coordinates s - zeta, delta = generate_longitudinal_coordinates(line=line, distribution='parabolic', + zeta, delta, matcher = generate_longitudinal_coordinates(line=line, distribution='parabolic', num_particles=num_particles, engine='single-rf-harmonic', sigma_z=sigma_z, - particle_ref=particle_ref) + particle_ref=particle_ref, return_matcher=True) # Initiate normalized coordinates x_norm = np.random.normal(size=num_particles) @@ -71,4 +72,7 @@ def parabolic_longitudinal_distribution(_context=None, nemitt_x=nemitt_x, nemitt_y=nemitt_y, weight=total_intensity_particles/num_particles, line=line) - return particles + if return_matcher: + return particles, matcher + else: + return particles From 92923038681cf7817ccb7954c291d5b87c76fef3 Mon Sep 17 00:00:00 2001 From: elwaagaa Date: Mon, 5 Jun 2023 10:01:33 +0200 Subject: [PATCH 7/8] Example to generate parabolic distribution. --- .../longitudinal/006_generate_parabolic.py | 76 +++++++++++++++++++ 1 file changed, 76 insertions(+) create mode 100644 examples/longitudinal/006_generate_parabolic.py diff --git a/examples/longitudinal/006_generate_parabolic.py b/examples/longitudinal/006_generate_parabolic.py new file mode 100644 index 0000000..9ccc836 --- /dev/null +++ b/examples/longitudinal/006_generate_parabolic.py @@ -0,0 +1,76 @@ +# copyright ############################### # +# This file is part of the Xpart Package. # +# Copyright (c) CERN, 2021. # +# ######################################### # + +""" +Simple example on how to generate a parabolic particle distribution +""" + +import json +import xpart as xp +import xtrack as xt +import xobjects as xo + +from xpart.longitudinal import parabolic_longitudinal_distribution +import matplotlib.pyplot as plt +import numpy as np + +# Load the reference particle +filename = xt._pkg_root.parent.joinpath('test_data/lhc_no_bb/line_and_particle.json') +with open(filename, 'r') as fid: + input_data = json.load(fid) +line = xt.Line.from_dict(input_data['line']) +line.build_tracker() + +# Specify the beam parameters +num_part = 1000000 +sigma_z = 0.05 + +# Build a reference particle +p0 = xp.Particles(mass0=xp.PROTON_MASS_EV, q0=1, p0c=7e12, x=1, y=3, + delta=[10]) + +# Built a set of three particles with different x coordinates +particles = parabolic_longitudinal_distribution( + num_particles=num_part, + nemitt_x=3e-6, + nemitt_y=3e-6, + sigma_z=sigma_z, + particle_ref=p0, + total_intensity_particles=1e10, + line=line + ) + + +# Make a histogram +bin_heights, bin_borders = np.histogram(particles.zeta, bins=300) +bin_widths = np.diff(bin_borders) +bin_centers = bin_borders[:-1] + bin_widths / 2 + + +# Generate the plots +plt.close('all') + +fig1 = plt.figure(1, figsize=(6.4, 7)) +ax21 = fig1.add_subplot(3,1,1) +ax22 = fig1.add_subplot(3,1,2) +ax23 = fig1.add_subplot(3,1,3) +ax21.plot(particles.x*1000, particles.px, '.', markersize=1) +ax21.set_xlabel(r'x [mm]') +ax21.set_ylabel(r'px [-]') +ax22.plot(particles.y*1000, particles.py, '.', markersize=1) +ax22.set_xlabel(r'y [mm]') +ax22.set_ylabel(r'py [-]') +ax23.plot(particles.zeta, particles.delta*1000, '.', markersize=1) +ax23.set_xlabel(r'z [-]') +ax23.set_ylabel(r'$\delta$ [1e-3]') +fig1.subplots_adjust(bottom=.08, top=.93, hspace=.33, left=.18, + right=.96, wspace=.33) + +fig2, ax2 = plt.subplots(1, 1, figsize = (6,5)) +ax2.bar(bin_centers, bin_heights, width=bin_widths) +ax2.set_ylabel('Counts') +ax2.set_xlabel(r'z [-]') + +plt.show() From f5a810f04d57e1e1509ec188456441fda0aac8cd Mon Sep 17 00:00:00 2001 From: elwaagaa Date: Mon, 5 Jun 2023 10:10:57 +0200 Subject: [PATCH 8/8] Corrected import --- examples/longitudinal/006_generate_parabolic.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/longitudinal/006_generate_parabolic.py b/examples/longitudinal/006_generate_parabolic.py index 9ccc836..8143610 100644 --- a/examples/longitudinal/006_generate_parabolic.py +++ b/examples/longitudinal/006_generate_parabolic.py @@ -12,7 +12,7 @@ import xtrack as xt import xobjects as xo -from xpart.longitudinal import parabolic_longitudinal_distribution +from xpart import parabolic_longitudinal_distribution import matplotlib.pyplot as plt import numpy as np