From 8a9c2724127cbc1d0fa48ce511a60870bed8ae95 Mon Sep 17 00:00:00 2001 From: Giovanni Iadarola Date: Sat, 13 Nov 2021 07:42:57 +0100 Subject: [PATCH 01/22] Moved linear_normal_form.py to xpart --- pymask/linear_normal_form.py | 223 ----------------------------------- 1 file changed, 223 deletions(-) delete mode 100644 pymask/linear_normal_form.py diff --git a/pymask/linear_normal_form.py b/pymask/linear_normal_form.py deleted file mode 100644 index 5cfdf74..0000000 --- a/pymask/linear_normal_form.py +++ /dev/null @@ -1,223 +0,0 @@ -import numpy as np -from scipy.optimize import fsolve - -import xline as xl -import xtrack as xt - -def _one_turn_map(p, particle_on_madx_co, tracker): - xl_part = particle_on_madx_co.copy() - xl_part.x = p[0] - xl_part.px = p[1] - xl_part.y = p[2] - xl_part.py = p[3] - xl_part.zeta = p[4] - xl_part.delta = p[5] - - part = xt.Particles(**xl_part.to_dict()) - tracker.track(part) - p_res = np.array([ - part.x[0], - part.px[0], - part.y[0], - part.py[0], - part.zeta[0], - part.delta[0]]) - return p_res - -def find_closed_orbit_from_tracker(tracker, particle_co_guess_dict): - particle_co_guess = xl.Particles.from_dict(particle_co_guess_dict) - print('Start CO search') - res = fsolve(lambda p: p - _one_turn_map(p, particle_co_guess, tracker), - x0=np.array([particle_co_guess.x, particle_co_guess.px, - particle_co_guess.y, particle_co_guess.py, - particle_co_guess.zeta, particle_co_guess.delta])) - print('Done CO search') - particle_on_co = particle_co_guess.copy() - particle_on_co.x = res[0] - particle_on_co.px = res[1] - particle_on_co.y = res[2] - particle_on_co.py = res[3] - particle_on_co.zeta = res[4] - particle_on_co.delta = res[5] - - return particle_on_co - -def compute_R_matrix_finite_differences( - particle_on_co, tracker, - dx=1e-9, dpx=1e-12, dy=1e-9, dpy=1e-12, - dzeta=1e-9, ddelta=1e-9, - symplectify=True): - - # Find R matrix - p0 = np.array([ - particle_on_co.x, - particle_on_co.px, - particle_on_co.y, - particle_on_co.py, - particle_on_co.zeta, - particle_on_co.delta]) - II = np.eye(6) - RR = np.zeros((6, 6), dtype=np.float64) - for jj, dd in enumerate([dx, dpx, dy, dpy, dzeta, ddelta]): - RR[:,jj]=(_one_turn_map(p0+II[jj]*dd, particle_on_co, tracker)- - _one_turn_map(p0-II[jj]*dd, particle_on_co, tracker))/(2*dd) - - if not 0.999 < np.linalg.det(RR) < 1.001: - raise ValueError('The determinant of the RR is out tolerance.') - - if symplectify: - return healy_symplectify(RR) - else: - return RR - -def healy_symplectify(M): - # https://accelconf.web.cern.ch/e06/PAPERS/WEPCH152.PDF - print("Symplectifying linear One-Turn-Map...") - - print("Before symplectifying: det(M) = {}".format(np.linalg.det(M))) - I = np.identity(6) - - S = np.array( - [ - [0.0, 1.0, 0.0, 0.0, 0.0, 0.0], - [-1.0, 0.0, 0.0, 0.0, 0.0, 0.0], - [0.0, 0.0, 0.0, 1.0, 0.0, 0.0], - [0.0, 0.0, -1.0, 0.0, 0.0, 0.0], - [0.0, 0.0, 0.0, 0.0, 0.0, 1.0], - [0.0, 0.0, 0.0, 0.0, -1.0, 0.0], - ] - ) - - V = np.matmul(S, np.matmul(I - M, np.linalg.inv(I + M))) - W = (V + V.T) / 2 - if np.linalg.det(I - np.matmul(S, W)) != 0: - M_new = np.matmul(I + np.matmul(S, W), - np.linalg.inv(I - np.matmul(S, W))) - else: - print("WARNING: det(I - SW) = 0!") - V_else = np.matmul(S, np.matmul(I + M, np.linalg.inv(I - M))) - W_else = (V_else + V_else.T) / 2 - M_new = -np.matmul( - I + np.matmul(S, W_else), np.linalg(I - np.matmul(S, W_else)) - ) - - print("After symplectifying: det(M) = {}".format(np.linalg.det(M_new))) - return M_new - -S = np.array([[0., 1., 0., 0., 0., 0.], - [-1., 0., 0., 0., 0., 0.], - [ 0., 0., 0., 1., 0., 0.], - [ 0., 0.,-1., 0., 0., 0.], - [ 0., 0., 0., 0., 0., 1.], - [ 0., 0., 0., 0.,-1., 0.]]) - -###################################################### -### Implement Normalization of fully coupled motion ## -###################################################### - -def Rot2D(mu): - return np.array([[ np.cos(mu), np.sin(mu)], - [-np.sin(mu), np.cos(mu)]]) - -def compute_linear_normal_form(M): - w0, v0 = np.linalg.eig(M) - - a0 = np.real(v0) - b0 = np.imag(v0) - - index_list = [0,5,1,2,3,4] - - ##### Sort modes in pairs of conjugate modes ##### - - conj_modes = np.zeros([3,2], dtype=np.int64) - for j in [0,1]: - conj_modes[j,0] = index_list[0] - del index_list[0] - - min_index = 0 - min_diff = abs(np.imag(w0[conj_modes[j,0]] + w0[index_list[min_index]])) - for i in range(1,len(index_list)): - diff = abs(np.imag(w0[conj_modes[j,0]] + w0[index_list[i]])) - if min_diff > diff: - min_diff = diff - min_index = i - - conj_modes[j,1] = index_list[min_index] - del index_list[min_index] - - conj_modes[2,0] = index_list[0] - conj_modes[2,1] = index_list[1] - - ################################################## - #### Select mode from pairs with positive (real @ S @ imag) ##### - - modes = np.empty(3, dtype=np.int64) - for ii,ind in enumerate(conj_modes): - if np.matmul(np.matmul(a0[:,ind[0]], S), b0[:,ind[0]]) > 0: - modes[ii] = ind[0] - else: - modes[ii] = ind[1] - - ################################################## - #### Sort modes such that (1,2,3) is close to (x,y,zeta) #### - for i in [1,2]: - if abs(v0[:,modes[0]])[0] < abs(v0[:,modes[i]])[0]: - modes[0], modes[i] = modes[i], modes[0] - - if abs(v0[:,modes[1]])[2] < abs(v0[:,modes[2]])[2]: - modes[2], modes[1] = modes[1], modes[2] - - ################################################## - #### Rotate eigenvectors to the Courant-Snyder parameterization #### - phase0 = np.log(v0[0,modes[0]]).imag - phase1 = np.log(v0[2,modes[1]]).imag - phase2 = np.log(v0[4,modes[2]]).imag - - v0[:,modes[0]] *= np.exp(-1.j*phase0) - v0[:,modes[1]] *= np.exp(-1.j*phase1) - v0[:,modes[2]] *= np.exp(-1.j*phase2) - - ################################################## - #### Construct W ################################# - - a1 = v0[:,modes[0]].real - a2 = v0[:,modes[1]].real - a3 = v0[:,modes[2]].real - b1 = v0[:,modes[0]].imag - b2 = v0[:,modes[1]].imag - b3 = v0[:,modes[2]].imag - - n1 = 1./np.sqrt(np.matmul(np.matmul(a1, S), b1)) - n2 = 1./np.sqrt(np.matmul(np.matmul(a2, S), b2)) - n3 = 1./np.sqrt(np.matmul(np.matmul(a3, S), b3)) - - a1 *= n1 - a2 *= n2 - a3 *= n3 - - b1 *= n1 - b2 *= n2 - b3 *= n3 - - W = np.array([a1,b1,a2,b2,a3,b3]).T - W[abs(W) < 1.e-14] = 0. # Set very small numbers to zero. - invW = np.matmul(np.matmul(S.T, W.T), S) - - ################################################## - #### Get tunes and rotation matrix in the normalized coordinates #### - - mu1 = np.log(w0[modes[0]]).imag - mu2 = np.log(w0[modes[1]]).imag - mu3 = np.log(w0[modes[2]]).imag - - #q1 = mu1/(2.*np.pi) - #q2 = mu2/(2.*np.pi) - #q3 = mu3/(2.*np.pi) - - R = np.zeros_like(W) - R[0:2,0:2] = Rot2D(mu1) - R[2:4,2:4] = Rot2D(mu2) - R[4:6,4:6] = Rot2D(mu3) - ################################################## - - return W, invW, R From ac5613990ad2f034a611cfa8c8c4bb035c6d89dc Mon Sep 17 00:00:00 2001 From: Giovanni Iadarola Date: Sat, 13 Nov 2021 08:44:20 +0100 Subject: [PATCH 02/22] Work in progress --- pymask/beambeam.py | 84 ++++++++++++++++++++++++++-------------------- 1 file changed, 47 insertions(+), 37 deletions(-) diff --git a/pymask/beambeam.py b/pymask/beambeam.py index 1bbf9d8..164976d 100644 --- a/pymask/beambeam.py +++ b/pymask/beambeam.py @@ -614,24 +614,16 @@ def find_bb_separations(points_weak, points_strong, names=None): return sep_x, sep_y - - - - - - - - def setup_beam_beam_in_line( line, bb_df, bb_coupling=False, ): - import xline + import xfields as xf assert bb_coupling is False # Not implemented - for ee, eename in zip(line.elements, line.element_names): - if isinstance(ee, xline.elements.BeamBeam4D): + for ii, (ee, eename) in enumerate(zip(line.elements, line.element_names)): + if isinstance(ee, xf.BeamBeamBiGaussian2D): ee.charge = (bb_df.loc[eename, 'other_num_particles'] * bb_df.loc[eename, 'other_particle_charge']) # TODO update xline interface to separate charge and b. population ee.sigma_x = np.sqrt(bb_df.loc[eename, 'other_Sigma_11']) @@ -640,31 +632,50 @@ def setup_beam_beam_in_line( ee.x_bb = bb_df.loc[eename, 'separation_x'] ee.y_bb = bb_df.loc[eename, 'separation_y'] - if isinstance(ee, xline.elements.BeamBeam6D): - - ee.phi = bb_df.loc[eename, 'phi'] - ee.alpha = bb_df.loc[eename, 'alpha'] - ee.x_bb_co = bb_df.loc[eename, 'separation_x'] - ee.y_bb_co = bb_df.loc[eename, 'separation_y'] - ee.charge_slices = [(bb_df.loc[eename, 'other_num_particles'] - * bb_df.loc[eename, 'other_particle_charge'])] # TODO update xline interface to separate charge and b. population - ee.zeta_slices = [0.0] - ee.sigma_11 = bb_df.loc[eename, 'other_Sigma_11'] - ee.sigma_12 = bb_df.loc[eename, 'other_Sigma_12'] - ee.sigma_13 = bb_df.loc[eename, 'other_Sigma_13'] - ee.sigma_14 = bb_df.loc[eename, 'other_Sigma_14'] - ee.sigma_22 = bb_df.loc[eename, 'other_Sigma_22'] - ee.sigma_23 = bb_df.loc[eename, 'other_Sigma_23'] - ee.sigma_24 = bb_df.loc[eename, 'other_Sigma_24'] - ee.sigma_33 = bb_df.loc[eename, 'other_Sigma_33'] - ee.sigma_34 = bb_df.loc[eename, 'other_Sigma_34'] - ee.sigma_44 = bb_df.loc[eename, 'other_Sigma_44'] + if isinstance(ee, xf.BeamBeamBiGaussian3D): + params = {} + params['phi'] = bb_df.loc[eename, 'phi'] + params['alpha'] = bb_df.loc[eename, 'alpha'] + params['x_bb_co'] = bb_df.loc[eename, 'separation_x'] + params['y_bb_co'] = bb_df.loc[eename, 'separation_y'] + # TODO update xline interface to separate charge and b. population + params['charge_slices'] = [(bb_df.loc[eename, 'other_num_particles'] + * bb_df.loc[eename, 'other_particle_charge'])] + params['zeta_slices'] = [0.0] + params['sigma_11'] = bb_df.loc[eename, 'other_Sigma_11'] + params['sigma_12'] = bb_df.loc[eename, 'other_Sigma_12'] + params['sigma_13'] = bb_df.loc[eename, 'other_Sigma_13'] + params['sigma_14'] = bb_df.loc[eename, 'other_Sigma_14'] + params['sigma_22'] = bb_df.loc[eename, 'other_Sigma_22'] + params['sigma_23'] = bb_df.loc[eename, 'other_Sigma_23'] + params['sigma_24'] = bb_df.loc[eename, 'other_Sigma_24'] + params['sigma_33'] = bb_df.loc[eename, 'other_Sigma_33'] + params['sigma_34'] = bb_df.loc[eename, 'other_Sigma_34'] + params['sigma_44'] = bb_df.loc[eename, 'other_Sigma_44'] if not (bb_coupling): - ee.sigma_13 = 0.0 - ee.sigma_14 = 0.0 - ee.sigma_23 = 0.0 - ee.sigma_24 = 0.0 + params['sigma_13'] = 0.0 + params['sigma_14'] = 0.0 + params['sigma_23'] = 0.0 + params['sigma_24'] = 0.0 + + params["x_co"] = 0 + params["px_co"] = 0 + params["y_co"] = 0 + params["py_co"] = 0 + params["zeta_co"] = 0 + params["delta_co"] = 0 + params["d_x"] = 0 + params["d_px"] = 0 + params["d_y"] = 0 + params["d_py"] = 0 + params["d_zeta"] = 0 + params["d_delta"] = 0 + + newee = xf.BeamBeamBiGaussian3D(old_interface=params) + line[ii] = newee + + def crabbing_strong_beam(mad, bb_dfs, z_crab_twiss, @@ -871,8 +882,7 @@ def find_bb_xma_yma(points_weak, points_strong, names=None): # Find as the position of the strong in the lab frame (points_strong[i_bb].p) # the reference frame of the weak in the lab frame (points_weak[i_bb].sp) - vbb_ws = points_strong[i_bb].p - points_weak[i_bb].sp - + vbb_ws = points_strong[i_bb].p - points_weak[i_bb].sp # Find separations xma.append(np.dot(vbb_ws, pbw.ex)) yma.append(np.dot(vbb_ws, pbw.ey)) @@ -880,7 +890,7 @@ def find_bb_xma_yma(points_weak, points_strong, names=None): return xma, yma def compute_xma_yma(bb_df): - + xma, yma = find_bb_xma_yma( points_weak=bb_df['self_lab_position'].values, points_strong=bb_df['other_lab_position'].values, From c72bbee83738a12b31ae7231258ec91df9f1ee9b Mon Sep 17 00:00:00 2001 From: Giovanni Iadarola Date: Sat, 13 Nov 2021 09:57:42 +0100 Subject: [PATCH 03/22] Adapted beambeam.py --- pymask/beambeam.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/pymask/beambeam.py b/pymask/beambeam.py index 164976d..0696bd2 100644 --- a/pymask/beambeam.py +++ b/pymask/beambeam.py @@ -624,14 +624,13 @@ def setup_beam_beam_in_line( for ii, (ee, eename) in enumerate(zip(line.elements, line.element_names)): if isinstance(ee, xf.BeamBeamBiGaussian2D): - ee.charge = (bb_df.loc[eename, 'other_num_particles'] - * bb_df.loc[eename, 'other_particle_charge']) # TODO update xline interface to separate charge and b. population + ee.n_particles=bb_df.loc[eename, 'other_num_particles'] + ee.q0 = bb_df.loc[eename, 'other_particle_charge']) ee.sigma_x = np.sqrt(bb_df.loc[eename, 'other_Sigma_11']) ee.sigma_y = np.sqrt(bb_df.loc[eename, 'other_Sigma_33']) - ee.beta_r = bb_df.loc[eename, 'other_relativistic_beta'] - ee.x_bb = bb_df.loc[eename, 'separation_x'] - ee.y_bb = bb_df.loc[eename, 'separation_y'] - + ee.beta0 = bb_df.loc[eename, 'other_relativistic_beta'] + ee.mean_x = bb_df.loc[eename, 'separation_x'] + ee.mean_y = bb_df.loc[eename, 'separation_y'] if isinstance(ee, xf.BeamBeamBiGaussian3D): params = {} params['phi'] = bb_df.loc[eename, 'phi'] From 9007a8aead2e2e37deefb53f0c486f797928af9e Mon Sep 17 00:00:00 2001 From: Giovanni Iadarola Date: Sat, 13 Nov 2021 14:21:17 +0100 Subject: [PATCH 04/22] Adapted to new xsuite, runs, not checked yet --- pymask/beambeam.py | 4 +- pymask/pymasktools.py | 52 ++++++++++--------- .../hl_lhc_collisions_python/000_pymask.py | 10 ++-- 3 files changed, 34 insertions(+), 32 deletions(-) diff --git a/pymask/beambeam.py b/pymask/beambeam.py index 0696bd2..633da55 100644 --- a/pymask/beambeam.py +++ b/pymask/beambeam.py @@ -625,7 +625,7 @@ def setup_beam_beam_in_line( for ii, (ee, eename) in enumerate(zip(line.elements, line.element_names)): if isinstance(ee, xf.BeamBeamBiGaussian2D): ee.n_particles=bb_df.loc[eename, 'other_num_particles'] - ee.q0 = bb_df.loc[eename, 'other_particle_charge']) + ee.q0 = bb_df.loc[eename, 'other_particle_charge'] ee.sigma_x = np.sqrt(bb_df.loc[eename, 'other_Sigma_11']) ee.sigma_y = np.sqrt(bb_df.loc[eename, 'other_Sigma_33']) ee.beta0 = bb_df.loc[eename, 'other_relativistic_beta'] @@ -672,7 +672,7 @@ def setup_beam_beam_in_line( params["d_delta"] = 0 newee = xf.BeamBeamBiGaussian3D(old_interface=params) - line[ii] = newee + line.elements[ii] = newee diff --git a/pymask/pymasktools.py b/pymask/pymasktools.py index 36bbc40..d5440f3 100644 --- a/pymask/pymasktools.py +++ b/pymask/pymasktools.py @@ -4,13 +4,11 @@ import numpy as np -import xline as xl import xtrack as xt +import xpart as xp +import xfields as xf from . import beambeam as bb -from .linear_normal_form import find_closed_orbit_from_tracker -from .linear_normal_form import compute_R_matrix_finite_differences -from .linear_normal_form import compute_linear_normal_form class JEncoder(json.JSONEncoder): def default(self, obj): @@ -210,7 +208,7 @@ def generate_sixtrack_input(mad, seq_name, bb_df, output_folder, ibbc_sixtrack, radius_sixtrack_multip_conversion_mad, skip_mad_use=False): - + six_fol_name = output_folder os.makedirs(six_fol_name, exist_ok=True) @@ -368,7 +366,7 @@ def get_optics_and_orbit_at_start_ring(mad, seq_name, with_bb_forces=False, mad_beam = mad.sequence[seq_name].beam assert mad_beam.deltap == 0, "Not implemented." - particle_on_madx_co = xl.Particles( + particle_on_madx_co = xp.Particles( p0c = mad_beam.pc*1e9, q0 = mad_beam.charge, mass0 = mad_beam.mass*1e9, @@ -474,25 +472,24 @@ def _set_orbit_dependent_parameters_for_bb(line, tracker, particle_on_co): ee.track(temp_particles) -def generate_xline(mad, seq_name, bb_df, +def generate_xtrack_line(mad, seq_name, bb_df, optics_and_co_at_start_ring_from_madx, folder_name=None, skip_mad_use=False, prepare_line_for_xtrack=True, steps_for_finite_diffs={'dx': 1e-9, 'dpx': 1e-12, 'dy': 1e-9, 'dpy': 1e-12, 'dzeta': 1e-9, 'ddelta': 1e-9}): - # Build xline model - print('Start building xline...') - line = xl.Line.from_madx_sequence( + # Build xsuite model + print('Start building xtrack line...') + line = xt.Line.from_madx_sequence( mad.sequence[seq_name], apply_madx_errors=True) - print('Done building xline.') + print('Done building xtrack.') if bb_df is not None: bb.setup_beam_beam_in_line(line, bb_df, bb_coupling=False) # Temporary fix due to bug in mad loader - cavities, cav_names = line.get_elements_of_type( - xl.elements.Cavity) + cavities, cav_names = line.get_elements_of_type(xt.Cavity) for cc, nn in zip(cavities, cav_names): if cc.frequency ==0.: ii_mad = mad.sequence[seq_name].element_names().index(nn) @@ -500,7 +497,7 @@ def generate_xline(mad, seq_name, bb_df, f0_mad = mad.sequence[seq_name].beam.freq0 * 1e6 # mad has it in MHz cc.frequency = f0_mad*cc_mad.parent.harmon - line_bb_dipole_not_cancelled_dict = line.to_dict(keepextra=True) + line_bb_dipole_not_cancelled_dict = line.to_dict() line_bb_dipole_not_cancelled_dict['particle_on_madx_co'] = ( optics_and_co_at_start_ring_from_madx['particle_on_madx_co']) line_bb_dipole_not_cancelled_dict['RR_madx'] = ( @@ -514,7 +511,7 @@ def generate_xline(mad, seq_name, bb_df, json.dump(line_bb_dipole_not_cancelled_dict, fid, cls=JEncoder) if prepare_line_for_xtrack: - tracker = xt.Tracker(sequence=line) + tracker = xt.Tracker(line=line) # Disable beam-beam for ee in tracker.line.elements: @@ -522,21 +519,26 @@ def generate_xline(mad, seq_name, bb_df, ee._temp_q0 = ee.q0 ee.q0 = 0 - particle_on_tracker_co = find_closed_orbit_from_tracker(tracker, - optics_and_co_at_start_ring_from_madx['particle_on_madx_co']) + particle_on_tracker_co = tracker.find_closed_orbit( + particle_co_guess=xp.Particles( + **optics_and_co_at_start_ring_from_madx['particle_on_madx_co'])) + + # Re-enable beam-beam + for ee in tracker.line.elements: + if ee.__class__.__name__.startswith('BeamBeam'): + ee.q0 = ee._temp_q0 + + xf.configure_orbit_dependent_parameters_for_bb(tracker, + particle_on_co=particle_on_tracker_co) - RR_finite_diffs = compute_R_matrix_finite_differences( - particle_on_tracker_co, tracker, symplectify=True, + RR_finite_diffs = tracker.compute_one_turn_matrix_finite_differences( + particle_on_tracker_co, **steps_for_finite_diffs) (WW_finite_diffs, WWInv_finite_diffs, RotMat_finite_diffs - ) = compute_linear_normal_form(RR_finite_diffs) - - # (Re-activates bb in line and tracker) - _set_orbit_dependent_parameters_for_bb(line, tracker, - particle_on_tracker_co) + ) = xp.compute_linear_normal_form(RR_finite_diffs) - line_bb_for_tracking_dict = line.to_dict(keepextra=True) + line_bb_for_tracking_dict = line.to_dict() line_bb_for_tracking_dict['particle_on_tracker_co'] = ( particle_on_tracker_co.to_dict()) line_bb_for_tracking_dict['RR_finite_diffs'] = RR_finite_diffs diff --git a/python_examples/hl_lhc_collisions_python/000_pymask.py b/python_examples/hl_lhc_collisions_python/000_pymask.py index 76dc68c..b6983b5 100644 --- a/python_examples/hl_lhc_collisions_python/000_pymask.py +++ b/python_examples/hl_lhc_collisions_python/000_pymask.py @@ -523,14 +523,14 @@ with open('./optics_orbit_at_start_ring_from_madx.json', 'w') as fid: json.dump(optics_and_co_at_start_ring_from_madx, fid, cls=pm.JEncoder) -################## -# Generate xline # -################## +######################## +# Generate xtrack line # +######################## if enable_bb_legacy: - print('xline is not generated with bb legacy macros') + print('xtrack line is not generated with bb legacy macros') else: - pm.generate_xline(mad_track, sequence_to_track, bb_df_track, + pm.generate_xtrack_line(mad_track, sequence_to_track, bb_df_track, optics_and_co_at_start_ring_from_madx, folder_name = './xlines', skip_mad_use=True, From 8b4f74c4d26f61b92255c5261532b90305017211 Mon Sep 17 00:00:00 2001 From: Giovanni Iadarola Date: Sat, 13 Nov 2021 14:24:14 +0100 Subject: [PATCH 05/22] Some cleanup --- pymask/pymasktools.py | 72 +------------------ .../hl_lhc_collisions_python/000_pymask.py | 4 +- 2 files changed, 3 insertions(+), 73 deletions(-) diff --git a/pymask/pymasktools.py b/pymask/pymasktools.py index d5440f3..c67936b 100644 --- a/pymask/pymasktools.py +++ b/pymask/pymasktools.py @@ -402,77 +402,7 @@ def get_optics_and_orbit_at_start_ring(mad, seq_name, with_bb_forces=False, - -def _set_orbit_dependent_parameters_for_bb(line, tracker, particle_on_co): - - temp_particles = xt.Particles(**particle_on_co.to_dict()) - for ii, ee in enumerate(tracker.line.elements): - if ee.__class__.__name__ == 'BeamBeamBiGaussian2D': - px_0 = temp_particles.px[0] - py_0 = temp_particles.py[0] - ee.q0 = ee._temp_q0 - - # Separation of 4D is so far set w.r.t. the closes orbit - # (to be able to compare against sixtrack) - # Here we set the righe quantities (coordinates of the strong beam) - ee.mean_x += temp_particles.x[0] - ee.mean_y += temp_particles.y[0] - line.elements[ii].x_bb = ee.mean_x - line.elements[ii].y_bb = ee.mean_y - - ee.track(temp_particles) - - ee.d_px = temp_particles.px - px_0 - ee.d_py = temp_particles.py - py_0 - line.elements[ii].d_px = ee.d_px - line.elements[ii].d_py = ee.d_py - - temp_particles.px -= ee.d_px - temp_particles.py -= ee.d_py - - elif ee.__class__.__name__ == 'BeamBeamBiGaussian3D': - ee.q0 = ee._temp_q0 - ee.x_CO = temp_particles.x[0] - ee.px_CO = temp_particles.px[0] - ee.y_CO = temp_particles.y[0] - ee.py_CO = temp_particles.py[0] - ee.sigma_CO = temp_particles.zeta[0] - ee.delta_CO = temp_particles.delta[0] - - ee.track(temp_particles) - - ee.Dx_sub = temp_particles.x[0] - ee.x_CO - ee.Dpx_sub = temp_particles.px[0] - ee.px_CO - ee.Dy_sub = temp_particles.y[0] - ee.y_CO - ee.Dpy_sub = temp_particles.py[0] - ee.py_CO - ee.Dsigma_sub = temp_particles.zeta[0] - ee.sigma_CO - ee.Ddelta_sub = temp_particles.delta[0] - ee.delta_CO - - temp_particles.x[0] = ee.x_CO - temp_particles.px[0] = ee.px_CO - temp_particles.y[0] = ee.y_CO - temp_particles.py[0] = ee.py_CO - temp_particles.zeta[0] = ee.sigma_CO - temp_particles.delta[0] = ee.delta_CO - - line.elements[ii].x_co = ee.x_CO - line.elements[ii].px_co = ee.px_CO - line.elements[ii].y_co = ee.y_CO - line.elements[ii].py_co = ee.py_CO - line.elements[ii].zeta_co = ee.sigma_CO - line.elements[ii].delta_co = ee.delta_CO - - line.elements[ii].d_x = ee.Dx_sub - line.elements[ii].d_px = ee.Dpx_sub - line.elements[ii].d_y = ee.Dy_sub - line.elements[ii].d_py = ee.Dpy_sub - line.elements[ii].d_zeta = ee.Dsigma_sub - line.elements[ii].d_delta = ee.Ddelta_sub - else: - ee.track(temp_particles) - - -def generate_xtrack_line(mad, seq_name, bb_df, +def generate_xsuite_line(mad, seq_name, bb_df, optics_and_co_at_start_ring_from_madx, folder_name=None, skip_mad_use=False, prepare_line_for_xtrack=True, diff --git a/python_examples/hl_lhc_collisions_python/000_pymask.py b/python_examples/hl_lhc_collisions_python/000_pymask.py index b6983b5..0b0c40b 100644 --- a/python_examples/hl_lhc_collisions_python/000_pymask.py +++ b/python_examples/hl_lhc_collisions_python/000_pymask.py @@ -530,9 +530,9 @@ if enable_bb_legacy: print('xtrack line is not generated with bb legacy macros') else: - pm.generate_xtrack_line(mad_track, sequence_to_track, bb_df_track, + pm.generate_xsuite_line(mad_track, sequence_to_track, bb_df_track, optics_and_co_at_start_ring_from_madx, - folder_name = './xlines', + folder_name = './xsuite_lines', skip_mad_use=True, prepare_line_for_xtrack=True) From 92407d07588d29349f17be53862b60bc790051ef Mon Sep 17 00:00:00 2001 From: Giovanni Iadarola Date: Sat, 13 Nov 2021 15:45:19 +0100 Subject: [PATCH 06/22] Working on comparison script (in progress...) --- .../checks_and_doc/t004_compare_xlines.py | 28 +++++++++---------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/python_examples/hl_lhc_collisions_python/checks_and_doc/t004_compare_xlines.py b/python_examples/hl_lhc_collisions_python/checks_and_doc/t004_compare_xlines.py index d6f366e..5c527f2 100644 --- a/python_examples/hl_lhc_collisions_python/checks_and_doc/t004_compare_xlines.py +++ b/python_examples/hl_lhc_collisions_python/checks_and_doc/t004_compare_xlines.py @@ -1,12 +1,14 @@ import time import shutil import pickle +import json import numpy as np -import xline +import xtrack as xt import sixtracktools + # Tests b1 with bb tests = [ { @@ -21,7 +23,7 @@ }, { 'test_name': 'B1 - pymask xline vs pymask sixtrack input', - 'path_test': '../xlines/line_bb_dipole_not_cancelled.json', + 'path_test': '../xsuite_lines/line_bb_dipole_not_cancelled.json', 'type_test': 'xline', 'path_ref': '../', 'type_ref': 'sixtrack', @@ -61,16 +63,13 @@ def norm(x): def prepare_line(path, input_type): if input_type == 'xline': - # Load xline machine - if path.endswith('.pkl'): - with open(path, 'rb') as fid: - ltest = xline.Line.from_dict(pickle.load(fid)) - else: - ltest = xline.Line.from_json(path) + # Load machine + with open(path, 'r') as fid: + ltest = xt.Line.from_dict(json.load(fid)) elif input_type == 'sixtrack': print('Build xline from sixtrack input:') sixinput_test = sixtracktools.sixinput.SixInput(path) - ltest = xline.Line.from_sixinput(sixinput_test) + ltest = xt.Line.from_sixinput(sixinput_test) print('Done') else: raise ValueError('What?!') @@ -135,8 +134,8 @@ def prepare_line(path, input_type): ): assert type(ee_test) == type(ee_six) - dtest = ee_test.to_dict(keepextra=True) - dref = ee_six.to_dict(keepextra=True) + dtest = ee_test.to_dict() + dref = ee_six.to_dict() for kk in dtest.keys(): @@ -144,6 +143,9 @@ def prepare_line(path, input_type): if np.isscalar(dref[kk]) and dtest[kk] == dref[kk]: continue + if isinstance(dref[kk], dict): + continue + # Check if the relative error is small val_test = dtest[kk] val_ref = dref[kk] @@ -172,9 +174,7 @@ def prepare_line(path, input_type): continue # Exception: drift length (100 um tolerance) - if not(strict) and isinstance( - ee_test, (xline.elements.Drift, xline.elements.DriftExact) - ): + if not(strict) and isinstance(ee_test, xt.Drift): if kk == "length": if diff_abs < 1e-4: continue From 0bee624d4cedb49a9fe13eeee7e7f60a95e6aec7 Mon Sep 17 00:00:00 2001 From: Giovanni Iadarola Date: Sat, 13 Nov 2021 22:23:12 +0100 Subject: [PATCH 07/22] Restored t004_compare_xlines.py --- .../checks_and_doc/t004_compare_xlines.py | 34 +++++++++++++------ 1 file changed, 24 insertions(+), 10 deletions(-) diff --git a/python_examples/hl_lhc_collisions_python/checks_and_doc/t004_compare_xlines.py b/python_examples/hl_lhc_collisions_python/checks_and_doc/t004_compare_xlines.py index 5c527f2..2b704bb 100644 --- a/python_examples/hl_lhc_collisions_python/checks_and_doc/t004_compare_xlines.py +++ b/python_examples/hl_lhc_collisions_python/checks_and_doc/t004_compare_xlines.py @@ -4,8 +4,10 @@ import json import numpy as np -import xtrack as xt import sixtracktools +import xtrack as xt +import xfields as xf + @@ -180,18 +182,24 @@ def prepare_line(path, input_type): continue # Exception: multipole lrad is not passed to sixtraxk - if isinstance(ee_test, xline.elements.Multipole): + if isinstance(ee_test, xt.Multipole): if kk == "length": if np.abs(ee_test.hxl) + np.abs(ee_test.hyl) == 0.0: continue - if kk == 'knl' or kk == 'ksl': + if kk == "order": + # Checked through bal + continue + if kk == 'knl' or kk == 'ksl' or kk == 'bal': if len(val_ref) != len(val_test): lmin = min(len(val_ref), len(val_test)) for vv in [val_ref,val_test]: if len(vv)> lmin: - for oo in range(lmin, len(vv)): # we do not care about errors above 10 - if vv[oo] != 0 and oo < 10: - raise ValueError('Missing significant multipole strength') + for oo in range(lmin, len(vv)): + # we do not care about errors above 10 + if vv[oo] != 0 and oo < {'knl':10, + 'ksl':10, 'bal':20}[kk]: + raise ValueError( + 'Missing significant multipole strength') val_ref = val_ref[:lmin] val_test = val_test[:lmin] @@ -226,7 +234,7 @@ def prepare_line(path, input_type): continue # Exceptions BB4D (separations are recalculated) - if not(strict) and isinstance(ee_test, xline.elements.BeamBeam4D): + if not(strict) and isinstance(ee_test, xf.BeamBeamBiGaussian2D): if kk == "x_bb": if diff_abs / dtest["sigma_x"] < 0.01: # This is neede to accommodate different leveling routines (1% difference) continue @@ -239,16 +247,22 @@ def prepare_line(path, input_type): if kk == "sigma_y": if diff_rel < 1e-5: continue + if isinstance(ee_test, xf.BeamBeamBiGaussian2D): + if kk == 'q0' or kk == 'n_particles': + # ambiguity due to old interface + if np.abs(ee_test.n_particles*ee_test.q0 - + ee_six.n_particles*ee_six.q0 ) < 1.: # charges + continue # Exceptions BB6D (angles and separations are recalculated) - if not(strict) and isinstance(ee_test, xline.elements.BeamBeam6D): + if not(strict) and isinstance(ee_test, xf.BeamBeamBiGaussian3D): if kk == "alpha": if diff_abs < 10e-6: continue - if kk == "x_co" or kk == "x_bb_co": + if kk == "x_co" or kk == "x_bb_co" or kk == 'delta_x': if diff_abs / np.sqrt(dtest["sigma_11"]) < 0.015: continue - if kk == "y_co" or kk == "y_bb_co": + if kk == "y_co" or kk == "y_bb_co" or kk == 'delta_y': if diff_abs / np.sqrt(dtest["sigma_33"]) < 0.015: continue if kk == "zeta_co": From 1d04c27523e2064794acb877ff9085594f09fbcf Mon Sep 17 00:00:00 2001 From: Giovanni Iadarola Date: Sat, 13 Nov 2021 22:56:24 +0100 Subject: [PATCH 08/22] Working on t005_check_crabs.py, there seems to be a discrepancy --- .../checks_and_doc/t005_check_crabs.py | 52 +++++++++---------- 1 file changed, 25 insertions(+), 27 deletions(-) diff --git a/python_examples/hl_lhc_collisions_python/checks_and_doc/t005_check_crabs.py b/python_examples/hl_lhc_collisions_python/checks_and_doc/t005_check_crabs.py index 8596a70..808bb58 100644 --- a/python_examples/hl_lhc_collisions_python/checks_and_doc/t005_check_crabs.py +++ b/python_examples/hl_lhc_collisions_python/checks_and_doc/t005_check_crabs.py @@ -2,7 +2,9 @@ import numpy as np import json -import xline +import xtrack as xt +import xpart as xp +import xfields as xf import sixtracktools L_lhc = 27e3 @@ -46,18 +48,19 @@ path_test = '../' type_test = 'sixtrack' -path_test = '../xlines/line_bb_dipole_not_cancelled.json' +path_test = '../xsuite_lines/line_bb_dipole_not_cancelled.json' type_test = 'xline' def prepare_line(path, input_type): if input_type == 'xline': # Load xline machine - ltest = xline.Line.from_json(path) + with open(path, 'r') as fid: + ltest = xt.Line.from_dict(json.load(fid)) elif input_type == 'sixtrack': print('Build xline from sixtrack input:') sixinput_test = sixtracktools.sixinput.SixInput(path) - ltest = xline.Line.from_sixinput(sixinput_test) + ltest = xt.Line.from_sixinput(sixinput_test) print('Done') else: raise ValueError('What?!') @@ -102,7 +105,7 @@ def prepare_line(path, input_type): R1_orbit = phi_weak * s_rel axcrab.plot(s_rel, np.array( - [getattr(bb, f'{plane}_bb_co') for bb in bb_elems])+R1_orbit, + [getattr(bb, f'delta_{plane}') for bb in bb_elems])+R1_orbit, 'o', color = 'r', alpha=.5, label='strong pyst') plt.plot(s_rel, R_no_crab + R_crab, '*', color='darkred', label='strong formula') @@ -119,12 +122,12 @@ def prepare_line(path, input_type): # Chek crabs in weak beam # Switch off all beam-beam lenses -bb_all, _ = ltest.get_elements_of_type([xline.elements.BeamBeam4D, - xline.elements.BeamBeam6D]) +bb_all, _ = ltest.get_elements_of_type([xf.BeamBeamBiGaussian2D, + xf.BeamBeamBiGaussian3D]) for bb in bb_all: bb.enabled = False # # Switch off all beam-beam lenses -crabs, crab_names = ltest.get_elements_of_type([xline.elements.RFMultipole]) +crabs, crab_names = ltest.get_elements_of_type([xt.RFMultipole]) #for cc in crabs: # cc.pn = [-90] # cc.ps = [-90] @@ -134,19 +137,13 @@ def prepare_line(path, input_type): with open('../optics_orbit_at_start_ring_from_madx.json', 'r') as fid: ddd = json.load(fid) -partco = xline.Particles.from_dict(ddd['particle_on_madx_co']) +partco = xp.Particles.from_dict(ddd['particle_on_madx_co']) z_slices = s_rel * 2.0 -partco.zeta += z_slices -partco.x += 0*z_slices -partco.s += 0*z_slices -partco.y += 0*z_slices -partco.px += 0*z_slices -partco.py += 0*z_slices -partco.delta += 0*z_slices +partco = xp.build_particles(particle_ref=partco, zeta=z_slices) -list_co = ltest.track_elem_by_elem(partco) -list_co2 = ltest.track_elem_by_elem(partco) +list_co = ltest.slow_track_elem_by_elem(partco) +list_co2 = ltest.slow_track_elem_by_elem(partco) plt.figure(2) axcox = plt.subplot(2,1,1) @@ -220,16 +217,17 @@ def prepare_line(path, input_type): axcbx.plot(s_twiss, x_twiss) axcby.plot(s_twiss, y_twiss) -part = xline.Particles.from_dict(ddd['particle_on_madx_co']) +part = xp.Particles.from_dict(ddd['particle_on_madx_co']) + z_test = np.array([0, z_crab_track]) -part.zeta += z_test -part.x += 0*z_test + np.array([0, x_twiss[0]]) -part.s += 0*z_test -part.y += 0*z_test + np.array([0, y_twiss[0]]) -part.px += 0*z_test + np.array([0, px_twiss[0]]) -part.py += 0*z_test + np.array([0, py_twiss[0]]) -part.delta += 0*z_test -list_track = ltest.track_elem_by_elem(part) +part = xp.build_particles(particle_ref=part, + zeta = z_test, + x = 0*z_test + np.array([0, x_twiss[0]]), + y = 0*z_test + np.array([0, y_twiss[0]]), + px = 0*z_test + np.array([0, px_twiss[0]]), + py = 0*z_test + np.array([0, py_twiss[0]]), + delta = 0*z_test) +list_track = ltest.slow_track_elem_by_elem(part) axcbx.plot([pp.s[0] for pp in list_track], [pp.x[1] - pp.x[0] for pp in list_track]) axcby.plot([pp.s[0] for pp in list_track], [pp.y[1] - pp.y[0] for pp in list_track]) From 0ff6c4f06dd8943864391a4fa63f8dccae515ce3 Mon Sep 17 00:00:00 2001 From: giadarol Date: Sun, 14 Nov 2021 20:34:25 +0100 Subject: [PATCH 09/22] Continue adapting examples --- .../checks_and_doc/t004_compare_xlines.py | 7 ++++++- .../checks_and_doc/t005_check_crabs.py | 7 ++++--- .../t007_check_orbit_and_lin_normal_form.py | 12 ++++++------ 3 files changed, 16 insertions(+), 10 deletions(-) diff --git a/python_examples/hl_lhc_collisions_python/checks_and_doc/t004_compare_xlines.py b/python_examples/hl_lhc_collisions_python/checks_and_doc/t004_compare_xlines.py index 2b704bb..57b51a2 100644 --- a/python_examples/hl_lhc_collisions_python/checks_and_doc/t004_compare_xlines.py +++ b/python_examples/hl_lhc_collisions_python/checks_and_doc/t004_compare_xlines.py @@ -146,7 +146,12 @@ def prepare_line(path, input_type): continue if isinstance(dref[kk], dict): - continue + if kk=='fieldmap': + continue + if kk=='boost_parameters': + continue + if kk=='Sigmas_0_star': + continue # Check if the relative error is small val_test = dtest[kk] diff --git a/python_examples/hl_lhc_collisions_python/checks_and_doc/t005_check_crabs.py b/python_examples/hl_lhc_collisions_python/checks_and_doc/t005_check_crabs.py index 808bb58..16f47c5 100644 --- a/python_examples/hl_lhc_collisions_python/checks_and_doc/t005_check_crabs.py +++ b/python_examples/hl_lhc_collisions_python/checks_and_doc/t005_check_crabs.py @@ -48,8 +48,8 @@ path_test = '../' type_test = 'sixtrack' -path_test = '../xsuite_lines/line_bb_dipole_not_cancelled.json' -type_test = 'xline' +#path_test = '../xsuite_lines/line_bb_dipole_not_cancelled.json' +#type_test = 'xline' def prepare_line(path, input_type): @@ -124,7 +124,8 @@ def prepare_line(path, input_type): # Switch off all beam-beam lenses bb_all, _ = ltest.get_elements_of_type([xf.BeamBeamBiGaussian2D, xf.BeamBeamBiGaussian3D]) -for bb in bb_all: bb.enabled = False +for bb in bb_all: + bb.q0 = 0 # # Switch off all beam-beam lenses crabs, crab_names = ltest.get_elements_of_type([xt.RFMultipole]) diff --git a/python_examples/hl_lhc_collisions_python/checks_and_doc/t007_check_orbit_and_lin_normal_form.py b/python_examples/hl_lhc_collisions_python/checks_and_doc/t007_check_orbit_and_lin_normal_form.py index bd74b32..fe5e1dd 100644 --- a/python_examples/hl_lhc_collisions_python/checks_and_doc/t007_check_orbit_and_lin_normal_form.py +++ b/python_examples/hl_lhc_collisions_python/checks_and_doc/t007_check_orbit_and_lin_normal_form.py @@ -2,18 +2,18 @@ import numpy as np -import xline as xl import xtrack as xt +import xpart as xp -with open('../xlines/line_bb_for_tracking.json', 'r') as fid: +with open('../xsuite_lines/line_bb_for_tracking.json', 'r') as fid: line_dict = json.load(fid) -line = xl.Line.from_dict(line_dict) +line = xt.Line.from_dict(line_dict) -partCO = xl.Particles.from_dict(line_dict['particle_on_tracker_co']) +partCO = xp.Particles.from_dict(line_dict['particle_on_tracker_co']) -tracker = xt.Tracker(sequence=line) -particles = xt.Particles(**partCO.to_dict()) +tracker = xt.Tracker(line=line) +particles = xp.Particles(**partCO.to_dict()) for _ in range(10): print(particles.at_turn[0], particles.x[0], particles.y[0], From c5d39df792b8b4e1056edf4af4c59700f9c2819a Mon Sep 17 00:00:00 2001 From: giadarol Date: Sun, 14 Nov 2021 22:17:26 +0100 Subject: [PATCH 10/22] Work in progress --- pymask/pymasktools.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/pymask/pymasktools.py b/pymask/pymasktools.py index c67936b..ef3295f 100644 --- a/pymask/pymasktools.py +++ b/pymask/pymasktools.py @@ -406,8 +406,8 @@ def generate_xsuite_line(mad, seq_name, bb_df, optics_and_co_at_start_ring_from_madx, folder_name=None, skip_mad_use=False, prepare_line_for_xtrack=True, - steps_for_finite_diffs={'dx': 1e-9, 'dpx': 1e-12, - 'dy': 1e-9, 'dpy': 1e-12, 'dzeta': 1e-9, 'ddelta': 1e-9}): + steps_for_finite_diffs={'dx': 1e-8, 'dpx': 1e-11, + 'dy': 1e-8, 'dpy': 1e-11, 'dzeta': 1e-7, 'ddelta': 1e-8}): # Build xsuite model print('Start building xtrack line...') @@ -461,10 +461,21 @@ def generate_xsuite_line(mad, seq_name, bb_df, xf.configure_orbit_dependent_parameters_for_bb(tracker, particle_on_co=particle_on_tracker_co) + # Disable beam-beam + for ee in tracker.line.elements: + if ee.__class__.__name__.startswith('BeamBeam'): + ee._temp_q0 = ee.q0 + ee.q0 = 0 + RR_finite_diffs = tracker.compute_one_turn_matrix_finite_differences( particle_on_tracker_co, **steps_for_finite_diffs) + # Re-enable beam-beam + for ee in tracker.line.elements: + if ee.__class__.__name__.startswith('BeamBeam'): + ee.q0 = ee._temp_q0 + (WW_finite_diffs, WWInv_finite_diffs, RotMat_finite_diffs ) = xp.compute_linear_normal_form(RR_finite_diffs) From d10adc2d23f8509c0a9787268b3920adcf7abf7c Mon Sep 17 00:00:00 2001 From: giadarol Date: Mon, 15 Nov 2021 13:51:19 +0100 Subject: [PATCH 11/22] Status --- pymask/pymasktools.py | 78 +++++++++++++++++++++++++++++++------------ 1 file changed, 57 insertions(+), 21 deletions(-) diff --git a/pymask/pymasktools.py b/pymask/pymasktools.py index ef3295f..0ba19f8 100644 --- a/pymask/pymasktools.py +++ b/pymask/pymasktools.py @@ -443,38 +443,21 @@ def generate_xsuite_line(mad, seq_name, bb_df, if prepare_line_for_xtrack: tracker = xt.Tracker(line=line) - # Disable beam-beam - for ee in tracker.line.elements: - if ee.__class__.__name__.startswith('BeamBeam'): - ee._temp_q0 = ee.q0 - ee.q0 = 0 - + _disable_beam_beam(tracker.line) particle_on_tracker_co = tracker.find_closed_orbit( particle_co_guess=xp.Particles( **optics_and_co_at_start_ring_from_madx['particle_on_madx_co'])) - - # Re-enable beam-beam - for ee in tracker.line.elements: - if ee.__class__.__name__.startswith('BeamBeam'): - ee.q0 = ee._temp_q0 + _restore_beam_beam(tracker.line) xf.configure_orbit_dependent_parameters_for_bb(tracker, particle_on_co=particle_on_tracker_co) - # Disable beam-beam - for ee in tracker.line.elements: - if ee.__class__.__name__.startswith('BeamBeam'): - ee._temp_q0 = ee.q0 - ee.q0 = 0 - + _disable_beam_beam(tracker.line) RR_finite_diffs = tracker.compute_one_turn_matrix_finite_differences( particle_on_tracker_co, **steps_for_finite_diffs) + _restore_beam_beam(tracker.line) - # Re-enable beam-beam - for ee in tracker.line.elements: - if ee.__class__.__name__.startswith('BeamBeam'): - ee.q0 = ee._temp_q0 (WW_finite_diffs, WWInv_finite_diffs, RotMat_finite_diffs ) = xp.compute_linear_normal_form(RR_finite_diffs) @@ -505,3 +488,56 @@ def save_mad_sequence_and_error(mad, seq_name, filename='lhc'): mad.select(flag="error",full=True) mad.esave(file=filename + "_errors_all.tfs") mad.save(sequence=seq_name,beam=True,file=filename + "_seq.madx") + + +def _disable_beam_beam(line): + for ee in line.elements: + if ee.__class__.__name__.startswith('BeamBeam'): + ee._temp_q0 = ee.q0 + ee.q0 = 0 + if ee.__class__.__name__ == 'BeamBeamBiGaussian2D': + ee._temp_d_px = ee.d_px + ee._temp_d_py = ee.d_py + ee.d_px = 0. + ee.d_py = 0. + elif ee.__class__.__name__ == 'BeamBeamBiGaussian3D': + ee._temp_Dx_sub = ee.Dx_sub + ee._temp_Dpx_sub = ee.Dpx_sub + ee._temp_Dy_sub = ee.Dy_sub + ee._temp_Dpy_sub = ee.Dpy_sub + ee._temp_Dsigma_sub = ee.Dsigma_sub + ee._temp_Ddelta_sub = ee.Ddelta_sub + ee.Dx_sub = 0. + ee.Dpx_sub = 0. + ee.Dy_sub = 0. + ee.Dpy_sub = 0. + ee.Dsigma_sub = 0. + ee.Ddelta_sub = 0. + else: + raise ValueError('What?!') + +def _restore_beam_beam(line): + for ee in line.elements: + if ee.__class__.__name__.startswith('BeamBeam'): + ee.q0 = ee._temp_q0 + del(ee._temp_q0) + if ee.__class__.__name__ == 'BeamBeamBiGaussian2D': + ee.d_px = ee._temp_d_px + ee.d_py = ee._temp_d_py + del(ee._temp_d_px) + del(ee._temp_d_py) + elif ee.__class__.__name__ == 'BeamBeamBiGaussian3D': + ee.Dx_sub = ee._temp_Dx_sub + ee.Dpx_sub = ee._temp_Dpx_sub + ee.Dy_sub = ee._temp_Dy_sub + ee.Dpy_sub = ee._temp_Dpy_sub + ee.Dsigma_sub = ee._temp_Dsigma_sub + ee.Ddelta_sub = ee._temp_Ddelta_sub + del(ee._temp_Dx_sub) + del(ee._temp_Dpx_sub) + del(ee._temp_Dy_sub) + del(ee._temp_Dpy_sub) + del(ee._temp_Dsigma_sub) + del(ee._temp_Ddelta_sub) + else: + raise ValueError('What?!') From 41985139bdd4b19376d2211cc75149828eef58c3 Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 16 Nov 2021 09:57:27 +0100 Subject: [PATCH 12/22] Adapting checks, work in progress --- pymask/pymasktools.py | 1 + .../hl_lhc_collisions_python/000_pymask.py | 4 ++-- .../t006_calc_and_plot_footprint.py | 13 ++++++----- .../t007_check_orbit_and_lin_normal_form.py | 22 +++---------------- 4 files changed, 13 insertions(+), 27 deletions(-) diff --git a/pymask/pymasktools.py b/pymask/pymasktools.py index 0ba19f8..14e1eda 100644 --- a/pymask/pymasktools.py +++ b/pymask/pymasktools.py @@ -449,6 +449,7 @@ def generate_xsuite_line(mad, seq_name, bb_df, **optics_and_co_at_start_ring_from_madx['particle_on_madx_co'])) _restore_beam_beam(tracker.line) + import pdb; pdb.set_trace() xf.configure_orbit_dependent_parameters_for_bb(tracker, particle_on_co=particle_on_tracker_co) diff --git a/python_examples/hl_lhc_collisions_python/000_pymask.py b/python_examples/hl_lhc_collisions_python/000_pymask.py index 0b0c40b..15de39c 100644 --- a/python_examples/hl_lhc_collisions_python/000_pymask.py +++ b/python_examples/hl_lhc_collisions_python/000_pymask.py @@ -194,7 +194,7 @@ print('Leveling not working in this mode!') else: if particle_type == 'ion': # the legacy macro for BB have been checked but not maintained - raise ValueError + raise ValueError # Luminosity levelling vars_for_legacy_level = ['lumi_ip8', 'nco_IP1', 'nco_IP2', 'nco_IP5', 'nco_IP8'] @@ -526,7 +526,7 @@ ######################## # Generate xtrack line # ######################## - +import pdb; pdb.set_trace() if enable_bb_legacy: print('xtrack line is not generated with bb legacy macros') else: diff --git a/python_examples/hl_lhc_collisions_python/checks_and_doc/t006_calc_and_plot_footprint.py b/python_examples/hl_lhc_collisions_python/checks_and_doc/t006_calc_and_plot_footprint.py index a77029d..54d033d 100644 --- a/python_examples/hl_lhc_collisions_python/checks_and_doc/t006_calc_and_plot_footprint.py +++ b/python_examples/hl_lhc_collisions_python/checks_and_doc/t006_calc_and_plot_footprint.py @@ -5,8 +5,9 @@ import footprint import matplotlib.pyplot as plt -import xline import sixtracktools +import xtrack as xt +import xpart as xp track_with = 'xtrack' #track_with = 'sixtrack' @@ -23,22 +24,22 @@ def prepare_line(path, input_type): if input_type == 'xline': # Load xline machine - ltest = xline.Line.from_json(path) + ltest = xt.Line.from_json(path) elif input_type == 'sixtrack': print('Build xline from sixtrack input:') sixinput_test = sixtracktools.sixinput.SixInput(path) - ltest = xline.Line.from_sixinput(sixinput_test) + ltest = xt.Line.from_sixinput(sixinput_test) print('Done') else: raise ValueError('What?!') return ltest -line = prepare_line('../xlines/line_bb_for_tracking.json', input_type='xline') +line = prepare_line('../xsuite_lines/line_bb_for_tracking.json', input_type='xline') -with open('../xlines/line_bb_for_tracking.json', 'r') as fid: - partCO = xline.Particles.from_dict( +with open('../xsuite_lines/line_bb_for_tracking.json', 'r') as fid: + partCO = xp.Particles.from_dict( json.load(fid)['particle_on_tracker_co']) with open('../optics_orbit_at_start_ring_from_madx.json', 'r') as fid: diff --git a/python_examples/hl_lhc_collisions_python/checks_and_doc/t007_check_orbit_and_lin_normal_form.py b/python_examples/hl_lhc_collisions_python/checks_and_doc/t007_check_orbit_and_lin_normal_form.py index fe5e1dd..a6b105d 100644 --- a/python_examples/hl_lhc_collisions_python/checks_and_doc/t007_check_orbit_and_lin_normal_form.py +++ b/python_examples/hl_lhc_collisions_python/checks_and_doc/t007_check_orbit_and_lin_normal_form.py @@ -22,7 +22,6 @@ WW = np.array(line_dict['WW_finite_diffs']) WWinv = np.array(line_dict['WWInv_finite_diffs']) - assert np.max(np.abs(np.dot(WW, WWinv) - np.eye(6)))<1e-10 @@ -35,24 +34,9 @@ x_norm = ampl_sigmas * np.sqrt(geom_emit_x) * np.cos(theta) px_norm = ampl_sigmas * np.sqrt(geom_emit_x) * np.sin(theta) -XX_norm= np.array([x_norm, - px_norm, - np.zeros(n_part), - np.zeros(n_part), - np.zeros(n_part), - np.zeros(n_part)]) - -XX = np.dot(WW, XX_norm) -pp = partCO.copy() - -pp.x += XX[0, :] -pp.px += XX[1, :] -pp.y += XX[2, :] -pp.py += XX[3, :] -pp.zeta += XX[4, :] -pp.delta += XX[5, :] -particles_matched = xt.Particles(**pp.to_dict()) - +particles_matched = xp.build_particles(particle_ref=partCO, + x_norm=x_norm, px_norm=px_norm, + R_matrix=np.array(line_dict['RR_finite_diffs'])) import matplotlib.pyplot as plt plt.close('all') From ccc69e4cd149e90c741db71f131b8b2fe2d315c9 Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 16 Nov 2021 10:33:28 +0100 Subject: [PATCH 13/22] Removed pdb --- pymask/pymasktools.py | 1 - python_examples/hl_lhc_collisions_python/000_pymask.py | 1 - 2 files changed, 2 deletions(-) diff --git a/pymask/pymasktools.py b/pymask/pymasktools.py index 14e1eda..0ba19f8 100644 --- a/pymask/pymasktools.py +++ b/pymask/pymasktools.py @@ -449,7 +449,6 @@ def generate_xsuite_line(mad, seq_name, bb_df, **optics_and_co_at_start_ring_from_madx['particle_on_madx_co'])) _restore_beam_beam(tracker.line) - import pdb; pdb.set_trace() xf.configure_orbit_dependent_parameters_for_bb(tracker, particle_on_co=particle_on_tracker_co) diff --git a/python_examples/hl_lhc_collisions_python/000_pymask.py b/python_examples/hl_lhc_collisions_python/000_pymask.py index 15de39c..fa009a2 100644 --- a/python_examples/hl_lhc_collisions_python/000_pymask.py +++ b/python_examples/hl_lhc_collisions_python/000_pymask.py @@ -526,7 +526,6 @@ ######################## # Generate xtrack line # ######################## -import pdb; pdb.set_trace() if enable_bb_legacy: print('xtrack line is not generated with bb legacy macros') else: From b9a2bb16310a8e8b1f933673feb3c30c6cdfa878 Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 16 Nov 2021 14:30:53 +0100 Subject: [PATCH 14/22] Fixed footprint --- .../checks_and_doc/helpers.py | 99 +++++-------------- .../t006_calc_and_plot_footprint.py | 26 ++--- 2 files changed, 33 insertions(+), 92 deletions(-) diff --git a/python_examples/hl_lhc_collisions_python/checks_and_doc/helpers.py b/python_examples/hl_lhc_collisions_python/checks_and_doc/helpers.py index 8caf3b0..e02a446 100644 --- a/python_examples/hl_lhc_collisions_python/checks_and_doc/helpers.py +++ b/python_examples/hl_lhc_collisions_python/checks_and_doc/helpers.py @@ -1,18 +1,19 @@ import numpy as np import os import sixtracktools -import xline import time +import xpart as xp +import xtrack as xt def vectorize_all_coords(Dx_wrt_CO_m, Dpx_wrt_CO_rad, Dy_wrt_CO_m, Dpy_wrt_CO_rad, - Dsigma_wrt_CO_m, Ddelta_wrt_CO): + Dzeta_wrt_CO_m, Ddelta_wrt_CO): n_part = 1 for var in [Dx_wrt_CO_m, Dpx_wrt_CO_rad, Dy_wrt_CO_m, Dpy_wrt_CO_rad, - Dsigma_wrt_CO_m, Ddelta_wrt_CO]: + Dzeta_wrt_CO_m, Ddelta_wrt_CO]: if hasattr(var, '__iter__'): if n_part == 1: n_part = len(var) @@ -22,26 +23,26 @@ def vectorize_all_coords(Dx_wrt_CO_m, Dpx_wrt_CO_rad, Dpx_wrt_CO_rad = Dpx_wrt_CO_rad + np.zeros(n_part) Dy_wrt_CO_m = Dy_wrt_CO_m + np.zeros(n_part) Dpy_wrt_CO_rad = Dpy_wrt_CO_rad + np.zeros(n_part) - Dsigma_wrt_CO_m = Dsigma_wrt_CO_m + np.zeros(n_part) + Dzeta_wrt_CO_m = Dzeta_wrt_CO_m + np.zeros(n_part) Ddelta_wrt_CO = Ddelta_wrt_CO + np.zeros(n_part) return Dx_wrt_CO_m, Dpx_wrt_CO_rad,\ Dy_wrt_CO_m, Dpy_wrt_CO_rad,\ - Dsigma_wrt_CO_m, Ddelta_wrt_CO + Dzeta_wrt_CO_m, Ddelta_wrt_CO def track_particle_sixtrack( partCO, Dx_wrt_CO_m, Dpx_wrt_CO_rad, Dy_wrt_CO_m, Dpy_wrt_CO_rad, - Dsigma_wrt_CO_m, Ddelta_wrt_CO, n_turns, + Dzeta_wrt_CO_m, Ddelta_wrt_CO, n_turns, input_folder='./'): Dx_wrt_CO_m, Dpx_wrt_CO_rad,\ Dy_wrt_CO_m, Dpy_wrt_CO_rad,\ - Dsigma_wrt_CO_m, Ddelta_wrt_CO = vectorize_all_coords( + Dzeta_wrt_CO_m, Ddelta_wrt_CO = vectorize_all_coords( Dx_wrt_CO_m, Dpx_wrt_CO_rad, Dy_wrt_CO_m, Dpy_wrt_CO_rad, - Dsigma_wrt_CO_m, Ddelta_wrt_CO) + Dzeta_wrt_CO_m, Ddelta_wrt_CO) n_part = len(Dx_wrt_CO_m) @@ -79,22 +80,22 @@ def track_particle_sixtrack( lines_f13 = [] for i_part in range(n_part): - if type(partCO) is xline.Particles: + if type(partCO) == xp.Particles: temp_part = partCO.copy() else: - temp_part = xline.Particles(**partCO) + temp_part = xp.Particles(**partCO) temp_part.x += Dx_wrt_CO_m[i_part] temp_part.px += Dpx_wrt_CO_rad[i_part] temp_part.y += Dy_wrt_CO_m[i_part] temp_part.py += Dpy_wrt_CO_rad[i_part] - temp_part.sigma += Dsigma_wrt_CO_m[i_part] - temp_part.delta += Ddelta_wrt_CO[i_part] + temp_part.zeta += Dzeta_wrt_CO_m[i_part] + temp_part.update_delta(temp_part.delta + Ddelta_wrt_CO[i_part]) lines_f13.append('%.10e\n' % ((temp_part.x) * 1e3)) lines_f13.append('%.10e\n' % ((temp_part.px) * temp_part.rpp * 1e3)) lines_f13.append('%.10e\n' % ((temp_part.y) * 1e3)) lines_f13.append('%.10e\n' % ((temp_part.py) * temp_part.rpp * 1e3)) - lines_f13.append('%.10e\n' % ((temp_part.sigma) * 1e3)) + lines_f13.append('%.10e\n' % ((temp_part.zeta) * 1e3)) lines_f13.append('%.10e\n' % ((temp_part.delta))) if i_part % 2 == 1: lines_f13.append('%.10e\n' % (temp_part.energy0 * 1e-6)) @@ -144,7 +145,7 @@ def track_particle_sixtrack( px_tbt = np.zeros((n_turns, n_part)) y_tbt = np.zeros((n_turns, n_part)) py_tbt = np.zeros((n_turns, n_part)) - sigma_tbt = np.zeros((n_turns, n_part)) + zeta_tbt = np.zeros((n_turns, n_part)) delta_tbt = np.zeros((n_turns, n_part)) for i_part in range(n_part): @@ -153,86 +154,36 @@ def track_particle_sixtrack( px_tbt[:, i_part] = sixdump_part.px y_tbt[:, i_part] = sixdump_part.y py_tbt[:, i_part] = sixdump_part.py - sigma_tbt[:, i_part] = sixdump_part.sigma + zeta_tbt[:, i_part] = sixdump_part.zeta delta_tbt[:, i_part] = sixdump_part.delta extra = {} - return x_tbt, px_tbt, y_tbt, py_tbt, sigma_tbt, delta_tbt, extra + return x_tbt, px_tbt, y_tbt, py_tbt, zeta_tbt, delta_tbt, extra -def track_particle_xline(line, part, Dx_wrt_CO_m, Dpx_wrt_CO_rad, - Dy_wrt_CO_m, Dpy_wrt_CO_rad, - Dsigma_wrt_CO_m, Ddelta_wrt_CO, n_turns, verbose=False): - - Dx_wrt_CO_m, Dpx_wrt_CO_rad,\ - Dy_wrt_CO_m, Dpy_wrt_CO_rad,\ - Dsigma_wrt_CO_m, Ddelta_wrt_CO = vectorize_all_coords( - Dx_wrt_CO_m, Dpx_wrt_CO_rad, - Dy_wrt_CO_m, Dpy_wrt_CO_rad, - Dsigma_wrt_CO_m, Ddelta_wrt_CO) - - part.x += Dx_wrt_CO_m - part.px += Dpx_wrt_CO_rad - part.y += Dy_wrt_CO_m - part.py += Dpy_wrt_CO_rad - part.sigma += Dsigma_wrt_CO_m - part.delta += Ddelta_wrt_CO - - x_tbt = [] - px_tbt = [] - y_tbt = [] - py_tbt = [] - sigma_tbt = [] - delta_tbt = [] - - for i_turn in range(n_turns): - if verbose: - print('Turn %d/%d' % (i_turn, n_turns)) - - x_tbt.append(part.x.copy()) - px_tbt.append(part.px.copy()) - y_tbt.append(part.y.copy()) - py_tbt.append(part.py.copy()) - sigma_tbt.append(part.sigma.copy()) - delta_tbt.append(part.delta.copy()) - - line.track(part) - - x_tbt = np.array(x_tbt) - px_tbt = np.array(px_tbt) - y_tbt = np.array(y_tbt) - py_tbt = np.array(py_tbt) - sigma_tbt = np.array(sigma_tbt) - delta_tbt = np.array(delta_tbt) - - extra = {} - - return x_tbt, px_tbt, y_tbt, py_tbt, sigma_tbt, delta_tbt def track_particle_xtrack( line, partCO, Dx_wrt_CO_m, Dpx_wrt_CO_rad, Dy_wrt_CO_m, Dpy_wrt_CO_rad, - Dsigma_wrt_CO_m, Ddelta_wrt_CO, n_turns, + Dzeta_wrt_CO_m, Ddelta_wrt_CO, n_turns, _context=None): - Dx_wrt_CO_m, Dpx_wrt_CO_rad,\ Dy_wrt_CO_m, Dpy_wrt_CO_rad,\ - Dsigma_wrt_CO_m, Ddelta_wrt_CO = vectorize_all_coords( + Dzeta_wrt_CO_m, Ddelta_wrt_CO = vectorize_all_coords( Dx_wrt_CO_m, Dpx_wrt_CO_rad, Dy_wrt_CO_m, Dpy_wrt_CO_rad, - Dsigma_wrt_CO_m, Ddelta_wrt_CO) + Dzeta_wrt_CO_m, Ddelta_wrt_CO) - import xtrack as xt - tracker = xt.Tracker(_context=_context, sequence=line) - particles = xt.Particles(_context=_context, + tracker = xt.Tracker(_context=_context, line=line) + particles = xp.Particles(_context=_context, p0c = partCO.p0c, x = partCO.x + Dx_wrt_CO_m, px = partCO.px + Dpx_wrt_CO_rad, y = partCO.y + Dy_wrt_CO_m, py = partCO.py + Dpy_wrt_CO_rad, - zeta = partCO.zeta + Dsigma_wrt_CO_m, + zeta = partCO.zeta + Dzeta_wrt_CO_m, delta = partCO.delta + Ddelta_wrt_CO) print('Start track') @@ -244,14 +195,14 @@ def track_particle_xtrack( px_tbt = tracker.record_last_track.px.copy().T y_tbt = tracker.record_last_track.y.copy().T py_tbt = tracker.record_last_track.py.copy().T - sigma_tbt = tracker.record_last_track.zeta.copy().T + zeta_tbt = tracker.record_last_track.zeta.copy().T delta_tbt = tracker.record_last_track.delta.copy().T print('Done loading!') extra = {'particles': particles, 'tracker': tracker} - return x_tbt, px_tbt, y_tbt, py_tbt, sigma_tbt, delta_tbt, extra + return x_tbt, px_tbt, y_tbt, py_tbt, zeta_tbt, delta_tbt, extra def betafun_from_ellip(x_tbt, px_tbt): diff --git a/python_examples/hl_lhc_collisions_python/checks_and_doc/t006_calc_and_plot_footprint.py b/python_examples/hl_lhc_collisions_python/checks_and_doc/t006_calc_and_plot_footprint.py index 54d033d..f74415d 100644 --- a/python_examples/hl_lhc_collisions_python/checks_and_doc/t006_calc_and_plot_footprint.py +++ b/python_examples/hl_lhc_collisions_python/checks_and_doc/t006_calc_and_plot_footprint.py @@ -10,7 +10,7 @@ import xpart as xp track_with = 'xtrack' -#track_with = 'sixtrack' +track_with = 'sixtrack' epsn_x = 2.5e-6 epsn_y = 2.5e-6 @@ -24,7 +24,8 @@ def prepare_line(path, input_type): if input_type == 'xline': # Load xline machine - ltest = xt.Line.from_json(path) + with open(path, 'r') as fid: + ltest = xt.Line.from_dict(json.load(fid)) elif input_type == 'sixtrack': print('Build xline from sixtrack input:') sixinput_test = sixtracktools.sixinput.SixInput(path) @@ -66,29 +67,18 @@ def prepare_line(path, input_type): DpxDpy_wrt_CO[ii, jj, 1] = xy_norm[ii, jj, 1] * np.sqrt(epsn_y / part.beta0 / part.gamma0 / beta_y) -if track_with == 'xline': - - part = partCO.copy() - - x_tbt, px_tbt, y_tbt, py_tbt, sigma_tbt, delta_tbt, extra = hp.track_particle_xline( - line, part=part, Dx_wrt_CO_m=0., Dpx_wrt_CO_rad=DpxDpy_wrt_CO[:, :, 0].flatten(), - Dy_wrt_CO_m=0, Dpy_wrt_CO_rad=DpxDpy_wrt_CO[:, :, 1].flatten(), - Dsigma_wrt_CO_m=0., Ddelta_wrt_CO=0., n_turns=n_turns, verbose=True) - - info = track_with - -elif track_with == 'sixtrack': - x_tbt, px_tbt, y_tbt, py_tbt, sigma_tbt, delta_tbt, extra = hp.track_particle_sixtrack( +if track_with == 'sixtrack': + x_tbt, px_tbt, y_tbt, py_tbt, zeta_tbt, delta_tbt, extra = hp.track_particle_sixtrack( partCO=partCO, Dx_wrt_CO_m=0., Dpx_wrt_CO_rad=DpxDpy_wrt_CO[:, :, 0].flatten(), Dy_wrt_CO_m=0, Dpy_wrt_CO_rad=DpxDpy_wrt_CO[:, :, 1].flatten(), - Dsigma_wrt_CO_m=0., Ddelta_wrt_CO=0., n_turns=n_turns, + Dzeta_wrt_CO_m=0., Ddelta_wrt_CO=0., n_turns=n_turns, input_folder='../') info = track_with elif track_with == 'xtrack': - x_tbt, px_tbt, y_tbt, py_tbt, sigma_tbt, delta_tbt, extra = hp.track_particle_xtrack( + x_tbt, px_tbt, y_tbt, py_tbt, zeta_tbt, delta_tbt, extra = hp.track_particle_xtrack( line=line, partCO=partCO, Dx_wrt_CO_m=0., Dpx_wrt_CO_rad=DpxDpy_wrt_CO[:, :, 0].flatten(), Dy_wrt_CO_m=0., Dpy_wrt_CO_rad=DpxDpy_wrt_CO[:, :, 1].flatten(), - Dsigma_wrt_CO_m=0., Ddelta_wrt_CO=0., n_turns=n_turns) + Dzeta_wrt_CO_m=0., Ddelta_wrt_CO=0., n_turns=n_turns) info = track_with else: raise ValueError('What?!') From 11dd3c2a5093914970f87a7f524ce7fce07519a7 Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 16 Nov 2021 14:42:42 +0100 Subject: [PATCH 15/22] Fixed last check script --- .../t008_check_against_sixtrack.py | 25 ++++++++----------- 1 file changed, 10 insertions(+), 15 deletions(-) diff --git a/python_examples/hl_lhc_collisions_python/checks_and_doc/t008_check_against_sixtrack.py b/python_examples/hl_lhc_collisions_python/checks_and_doc/t008_check_against_sixtrack.py index 490c01c..d443a0a 100644 --- a/python_examples/hl_lhc_collisions_python/checks_and_doc/t008_check_against_sixtrack.py +++ b/python_examples/hl_lhc_collisions_python/checks_and_doc/t008_check_against_sixtrack.py @@ -2,11 +2,10 @@ import numpy as np import NAFFlib import helpers as hp -import footprint import matplotlib.pyplot as plt -import xline as xl import xtrack as xt +import xpart as xp import sixtracktools @@ -15,33 +14,29 @@ num_turns = 5 -with open('../xlines/line_bb_for_tracking.json', 'r') as fid: +with open('../xsuite_lines/line_bb_for_tracking.json', 'r') as fid: dict_line_xtrack = json.load(fid) #with open('../xline/line_bb_dipole_cancelled.json', 'r') as fid: # dict_line_old = json.load(fid) -line = xl.Line.from_dict(dict_line_xtrack) -#line = xl.Line.from_dict(dict_line_old) +line = xt.Line.from_dict(dict_line_xtrack) -partCO = xl.Particles.from_dict(dict_line_xtrack['particle_on_tracker_co']) +partCO = xp.Particles.from_dict(dict_line_xtrack['particle_on_tracker_co']) (x_tbt_sixtrack, px_tbt_sixtrack, y_tbt_sixtrack, py_tbt_sixtrack, - sigma_tbt_sixtrack, delta_tbt_sixtrack, extra) = hp.track_particle_sixtrack( + zeta_tbt_sixtrack, delta_tbt_sixtrack, extra) = hp.track_particle_sixtrack( partCO=partCO, Dx_wrt_CO_m=np.array(displace_x), Dpx_wrt_CO_rad=0, Dy_wrt_CO_m=np.array(displace_y), Dpy_wrt_CO_rad=0., - Dsigma_wrt_CO_m=0., Ddelta_wrt_CO=0., n_turns=num_turns, + Dzeta_wrt_CO_m=0., Ddelta_wrt_CO=0., n_turns=num_turns, input_folder='../') -tracker = xt.Tracker(sequence=line) +tracker = xt.Tracker(line=line) -part_track = partCO.copy() -part_track.x += np.array(displace_x) -part_track.y += np.array(displace_y) - -particles = xt.Particles(**part_track.to_dict()) -particles.particle_id = np.arange(particles.num_particles) +particles = xp.build_particles(particle_ref=partCO, + x=np.array(displace_x), + y=np.array(displace_y)) tracker.track(particles, turn_by_turn_monitor=True, num_turns=num_turns) From 58a6e1cf7a2e82de1491d56521e1c0317fdda14c Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 16 Nov 2021 14:45:26 +0100 Subject: [PATCH 16/22] Rename --- .../{t004_compare_xlines.py => t004_check_output_consistency.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename python_examples/hl_lhc_collisions_python/checks_and_doc/{t004_compare_xlines.py => t004_check_output_consistency.py} (100%) diff --git a/python_examples/hl_lhc_collisions_python/checks_and_doc/t004_compare_xlines.py b/python_examples/hl_lhc_collisions_python/checks_and_doc/t004_check_output_consistency.py similarity index 100% rename from python_examples/hl_lhc_collisions_python/checks_and_doc/t004_compare_xlines.py rename to python_examples/hl_lhc_collisions_python/checks_and_doc/t004_check_output_consistency.py From 2940731b7bba7e7869728baa2134452cfaefeabb Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 16 Nov 2021 14:54:02 +0100 Subject: [PATCH 17/22] Strengthening tests on orbit and optics --- .../t007_check_orbit_and_lin_normal_form.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/python_examples/hl_lhc_collisions_python/checks_and_doc/t007_check_orbit_and_lin_normal_form.py b/python_examples/hl_lhc_collisions_python/checks_and_doc/t007_check_orbit_and_lin_normal_form.py index a6b105d..676aae7 100644 --- a/python_examples/hl_lhc_collisions_python/checks_and_doc/t007_check_orbit_and_lin_normal_form.py +++ b/python_examples/hl_lhc_collisions_python/checks_and_doc/t007_check_orbit_and_lin_normal_form.py @@ -25,7 +25,7 @@ assert np.max(np.abs(np.dot(WW, WWinv) - np.eye(6)))<1e-10 -ampl_sigmas = 1. +ampl_sigmas = 0.2 norm_emit_x = 2.5e-6 geom_emit_x = norm_emit_x / particles.beta0 / particles.gamma0 @@ -37,14 +37,13 @@ particles_matched = xp.build_particles(particle_ref=partCO, x_norm=x_norm, px_norm=px_norm, R_matrix=np.array(line_dict['RR_finite_diffs'])) +particles = particles_matched.copy() +tracker.track(particles_matched, num_turns=10) import matplotlib.pyplot as plt plt.close('all') plt.figure(1) plt.plot(particles_matched.x, particles_matched.px) - -tracker.track(particles_matched, num_turns=10) - -plt.plot(particles_matched.x, particles_matched.px) +plt.plot(particles.x, particles.px) plt.show() From 531e5c1f1791e2d64af124960175b872b5f275de Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 16 Nov 2021 16:27:50 +0100 Subject: [PATCH 18/22] Some automatic checks --- .../t007_check_orbit_and_lin_normal_form.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/python_examples/hl_lhc_collisions_python/checks_and_doc/t007_check_orbit_and_lin_normal_form.py b/python_examples/hl_lhc_collisions_python/checks_and_doc/t007_check_orbit_and_lin_normal_form.py index 676aae7..25c6fec 100644 --- a/python_examples/hl_lhc_collisions_python/checks_and_doc/t007_check_orbit_and_lin_normal_form.py +++ b/python_examples/hl_lhc_collisions_python/checks_and_doc/t007_check_orbit_and_lin_normal_form.py @@ -20,6 +20,9 @@ particles.zeta[0]) tracker.track(particles) +for nn in 'x px y py zeta delta'.split(): + assert np.abs(getattr(particles, nn) - getattr(partCO, nn)) < 1e-11 + WW = np.array(line_dict['WW_finite_diffs']) WWinv = np.array(line_dict['WWInv_finite_diffs']) assert np.max(np.abs(np.dot(WW, WWinv) - np.eye(6)))<1e-10 @@ -37,13 +40,19 @@ particles_matched = xp.build_particles(particle_ref=partCO, x_norm=x_norm, px_norm=px_norm, R_matrix=np.array(line_dict['RR_finite_diffs'])) -particles = particles_matched.copy() -tracker.track(particles_matched, num_turns=10) +particles_test = particles_matched.copy() +tracker.track(particles_test, num_turns=10) + +i_matched = np.argmax(particles_matched.x) +i_test = np.argmax(particles_test.x) +assert np.abs(particles_test.x[i_test] - particles_matched.x[i_matched]) < 1e-6 +assert np.abs(particles_test.px[i_test] - particles_matched.px[i_matched]) < 1e-7 + import matplotlib.pyplot as plt plt.close('all') plt.figure(1) plt.plot(particles_matched.x, particles_matched.px) -plt.plot(particles.x, particles.px) +plt.plot(particles_test.x, particles_test.px) plt.show() From 2c01b37dda51743e22262ab1df433a75f491571f Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 16 Nov 2021 17:10:53 +0100 Subject: [PATCH 19/22] Some cleanup --- pymask/beambeam.py | 2 +- .../hl_lhc_collisions_python/001_reload.py | 1 - .../t004_check_output_consistency.py | 14 +++++++------- .../checks_and_doc/t005_check_crabs.py | 12 ++++++------ .../checks_and_doc/t006_calc_and_plot_footprint.py | 8 ++++---- .../checks_and_doc/t008_check_against_sixtrack.py | 3 --- .../checks_and_doc/tests.md | 12 ++++++------ .../t002_test_save_lines.py | 5 ++--- 8 files changed, 26 insertions(+), 31 deletions(-) diff --git a/pymask/beambeam.py b/pymask/beambeam.py index 633da55..efe8024 100644 --- a/pymask/beambeam.py +++ b/pymask/beambeam.py @@ -637,7 +637,7 @@ def setup_beam_beam_in_line( params['alpha'] = bb_df.loc[eename, 'alpha'] params['x_bb_co'] = bb_df.loc[eename, 'separation_x'] params['y_bb_co'] = bb_df.loc[eename, 'separation_y'] - # TODO update xline interface to separate charge and b. population + # TODO update xtrack interface to separate charge and b. population params['charge_slices'] = [(bb_df.loc[eename, 'other_num_particles'] * bb_df.loc[eename, 'other_particle_charge'])] params['zeta_slices'] = [0.0] diff --git a/python_examples/hl_lhc_collisions_python/001_reload.py b/python_examples/hl_lhc_collisions_python/001_reload.py index b5c7f6a..7e6ec86 100644 --- a/python_examples/hl_lhc_collisions_python/001_reload.py +++ b/python_examples/hl_lhc_collisions_python/001_reload.py @@ -1,6 +1,5 @@ import sys import numpy as np -import xline as xl import xtrack as xt sys.path.append('./modules') diff --git a/python_examples/hl_lhc_collisions_python/checks_and_doc/t004_check_output_consistency.py b/python_examples/hl_lhc_collisions_python/checks_and_doc/t004_check_output_consistency.py index 57b51a2..3d2d8c6 100644 --- a/python_examples/hl_lhc_collisions_python/checks_and_doc/t004_check_output_consistency.py +++ b/python_examples/hl_lhc_collisions_python/checks_and_doc/t004_check_output_consistency.py @@ -24,9 +24,9 @@ 'strict': False, }, { - 'test_name': 'B1 - pymask xline vs pymask sixtrack input', + 'test_name': 'B1 - pymask xsuite vs pymask sixtrack input', 'path_test': '../xsuite_lines/line_bb_dipole_not_cancelled.json', - 'type_test': 'xline', + 'type_test': 'xsuite', 'path_ref': '../', 'type_ref': 'sixtrack', 'rtol': 4e-7, @@ -48,9 +48,9 @@ # 'strict': False, # }, # { -# 'test_name': 'B4 - pymask xline vs pymask sixtrack input', -# 'path_test': '../xline/line_bb_dipole_not_cancelled.json', -# 'type_test': 'xline', +# 'test_name': 'B4 - pymask xsuite vs pymask sixtrack input', +# 'path_test': '../xsuite/line_bb_dipole_not_cancelled.json', +# 'type_test': 'xsuite', # 'path_ref': '../', # 'type_ref': 'sixtrack', # 'rtol': 4e-7, @@ -64,12 +64,12 @@ def norm(x): def prepare_line(path, input_type): - if input_type == 'xline': + if input_type == 'xsuite': # Load machine with open(path, 'r') as fid: ltest = xt.Line.from_dict(json.load(fid)) elif input_type == 'sixtrack': - print('Build xline from sixtrack input:') + print('Build xsuite from sixtrack input:') sixinput_test = sixtracktools.sixinput.SixInput(path) ltest = xt.Line.from_sixinput(sixinput_test) print('Done') diff --git a/python_examples/hl_lhc_collisions_python/checks_and_doc/t005_check_crabs.py b/python_examples/hl_lhc_collisions_python/checks_and_doc/t005_check_crabs.py index 16f47c5..835f5db 100644 --- a/python_examples/hl_lhc_collisions_python/checks_and_doc/t005_check_crabs.py +++ b/python_examples/hl_lhc_collisions_python/checks_and_doc/t005_check_crabs.py @@ -49,16 +49,16 @@ type_test = 'sixtrack' #path_test = '../xsuite_lines/line_bb_dipole_not_cancelled.json' -#type_test = 'xline' +#type_test = 'xsuite' def prepare_line(path, input_type): - if input_type == 'xline': - # Load xline machine + if input_type == 'xsuite': + # Load xsuite machine with open(path, 'r') as fid: ltest = xt.Line.from_dict(json.load(fid)) elif input_type == 'sixtrack': - print('Build xline from sixtrack input:') + print('Build xsuite from sixtrack input:') sixinput_test = sixtracktools.sixinput.SixInput(path) ltest = xt.Line.from_sixinput(sixinput_test) print('Done') @@ -106,7 +106,7 @@ def prepare_line(path, input_type): axcrab.plot(s_rel, np.array( [getattr(bb, f'delta_{plane}') for bb in bb_elems])+R1_orbit, - 'o', color = 'r', alpha=.5, label='strong pyst') + 'o', color = 'r', alpha=.5, label='strong xsuite') plt.plot(s_rel, R_no_crab + R_crab, '*', color='darkred', label='strong formula') #axcrab.plot(s_rel, np.array([bb.y_bb_co for bb in bb_elems]), 'o', color='r', alpha=.5) @@ -180,7 +180,7 @@ def prepare_line(path, input_type): for ibb, bb in enumerate(bb_elems): r_lenses.append(getattr(list_co[bb_index[ibb]], plane)[ibb]) -axcrab.plot(s_rel, r_lenses, 'o', color='b', alpha=.5, label= 'weak xline') +axcrab.plot(s_rel, r_lenses, 'o', color='b', alpha=.5, label= 'weak xsuite') Rw_crab = phi_c_weak * L_lhc / (2*np.pi*h_cc) *np.sin(2*np.pi*h_cc/L_lhc*2*s_rel) Rw_no_crab = phi_weak * s_rel axcrab.plot(s_rel, Rw_crab + Rw_no_crab, '*', color='darkblue', diff --git a/python_examples/hl_lhc_collisions_python/checks_and_doc/t006_calc_and_plot_footprint.py b/python_examples/hl_lhc_collisions_python/checks_and_doc/t006_calc_and_plot_footprint.py index f74415d..90aa300 100644 --- a/python_examples/hl_lhc_collisions_python/checks_and_doc/t006_calc_and_plot_footprint.py +++ b/python_examples/hl_lhc_collisions_python/checks_and_doc/t006_calc_and_plot_footprint.py @@ -22,12 +22,12 @@ def prepare_line(path, input_type): - if input_type == 'xline': - # Load xline machine + if input_type == 'xsuite': + # Load xsuite machine with open(path, 'r') as fid: ltest = xt.Line.from_dict(json.load(fid)) elif input_type == 'sixtrack': - print('Build xline from sixtrack input:') + print('Build xsuite from sixtrack input:') sixinput_test = sixtracktools.sixinput.SixInput(path) ltest = xt.Line.from_sixinput(sixinput_test) print('Done') @@ -36,7 +36,7 @@ def prepare_line(path, input_type): return ltest -line = prepare_line('../xsuite_lines/line_bb_for_tracking.json', input_type='xline') +line = prepare_line('../xsuite_lines/line_bb_for_tracking.json', input_type='xsuite') with open('../xsuite_lines/line_bb_for_tracking.json', 'r') as fid: diff --git a/python_examples/hl_lhc_collisions_python/checks_and_doc/t008_check_against_sixtrack.py b/python_examples/hl_lhc_collisions_python/checks_and_doc/t008_check_against_sixtrack.py index d443a0a..5689b68 100644 --- a/python_examples/hl_lhc_collisions_python/checks_and_doc/t008_check_against_sixtrack.py +++ b/python_examples/hl_lhc_collisions_python/checks_and_doc/t008_check_against_sixtrack.py @@ -17,9 +17,6 @@ with open('../xsuite_lines/line_bb_for_tracking.json', 'r') as fid: dict_line_xtrack = json.load(fid) -#with open('../xline/line_bb_dipole_cancelled.json', 'r') as fid: -# dict_line_old = json.load(fid) - line = xt.Line.from_dict(dict_line_xtrack) partCO = xp.Particles.from_dict(dict_line_xtrack['particle_on_tracker_co']) diff --git a/python_examples/hl_lhc_collisions_python/checks_and_doc/tests.md b/python_examples/hl_lhc_collisions_python/checks_and_doc/tests.md index 6e97704..b000588 100644 --- a/python_examples/hl_lhc_collisions_python/checks_and_doc/tests.md +++ b/python_examples/hl_lhc_collisions_python/checks_and_doc/tests.md @@ -52,7 +52,7 @@ We execute the reference (can be reused from the previous test): python ../../unmask.py main.mask parameters_for_unmask.txt --run-cpymad ``` In the folder ```checks_and_doc```: -Configure ```t004_compare_xline_lines.py``` for the test: +Configure ```t004_check_output_consistency.py``` for the test: ```python # Test b1 tests = [ @@ -67,9 +67,9 @@ tests = [ 'strict': False, }, { - 'test_name': 'B1 - pymask xline vs pymask sixtrack input', - 'path_test': '../xline/line_bb_dipole_not_cancelled.json', - 'type_test': 'xline', + 'test_name': 'B1 - pymask xsuite vs pymask sixtrack input', + 'path_test': '../xsuite/line_bb_dipole_not_cancelled.json', + 'type_test': 'xsuite', 'path_ref': '../', 'type_ref': 'sixtrack', 'rtol': 4e-7, @@ -79,10 +79,10 @@ tests = [ ] ``` -Check using xline lines: +Check using xsuite lines: ```bash python t003_fc_to_fort.py -python t004_compare_xlines.py.py +python t004_check_output_consistency.py ``` **Check crab cavities:** diff --git a/python_examples/hl_lhc_collisions_python/t002_test_save_lines.py b/python_examples/hl_lhc_collisions_python/t002_test_save_lines.py index 21e6f32..bae12f9 100644 --- a/python_examples/hl_lhc_collisions_python/t002_test_save_lines.py +++ b/python_examples/hl_lhc_collisions_python/t002_test_save_lines.py @@ -5,7 +5,6 @@ import numpy as np -import xline as xl import xtrack as xt from scipy.optimize import fsolve @@ -33,8 +32,8 @@ with open('./optics_orbit_at_start_ring_from_madx.json', 'r') as fid: optics_and_co_at_start_ring_from_madx = json.load(fid) -pm.generate_xline(mad, seq_name, bb_df_track, +pm.generate_xsuite_line(mad, seq_name, bb_df_track, optics_and_co_at_start_ring_from_madx, - folder_name = 'xlines', + folder_name = 'xsuite_lines', skip_mad_use=False, prepare_line_for_xtrack=True) From 2dd071ab8f70b55b4eef39c829b9c3419d891c6d Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 16 Nov 2021 18:02:18 +0100 Subject: [PATCH 20/22] Cleanup --- python_examples/ions_python/000_pymask.py | 15 ++- .../ions_python/checks_and_doc/helpers.py | 99 +++++-------------- .../t008_check_against_sixtrack.py | 30 +++--- 3 files changed, 44 insertions(+), 100 deletions(-) diff --git a/python_examples/ions_python/000_pymask.py b/python_examples/ions_python/000_pymask.py index 76dc68c..fa009a2 100644 --- a/python_examples/ions_python/000_pymask.py +++ b/python_examples/ions_python/000_pymask.py @@ -194,7 +194,7 @@ print('Leveling not working in this mode!') else: if particle_type == 'ion': # the legacy macro for BB have been checked but not maintained - raise ValueError + raise ValueError # Luminosity levelling vars_for_legacy_level = ['lumi_ip8', 'nco_IP1', 'nco_IP2', 'nco_IP5', 'nco_IP8'] @@ -523,16 +523,15 @@ with open('./optics_orbit_at_start_ring_from_madx.json', 'w') as fid: json.dump(optics_and_co_at_start_ring_from_madx, fid, cls=pm.JEncoder) -################## -# Generate xline # -################## - +######################## +# Generate xtrack line # +######################## if enable_bb_legacy: - print('xline is not generated with bb legacy macros') + print('xtrack line is not generated with bb legacy macros') else: - pm.generate_xline(mad_track, sequence_to_track, bb_df_track, + pm.generate_xsuite_line(mad_track, sequence_to_track, bb_df_track, optics_and_co_at_start_ring_from_madx, - folder_name = './xlines', + folder_name = './xsuite_lines', skip_mad_use=True, prepare_line_for_xtrack=True) diff --git a/python_examples/ions_python/checks_and_doc/helpers.py b/python_examples/ions_python/checks_and_doc/helpers.py index 8caf3b0..e02a446 100644 --- a/python_examples/ions_python/checks_and_doc/helpers.py +++ b/python_examples/ions_python/checks_and_doc/helpers.py @@ -1,18 +1,19 @@ import numpy as np import os import sixtracktools -import xline import time +import xpart as xp +import xtrack as xt def vectorize_all_coords(Dx_wrt_CO_m, Dpx_wrt_CO_rad, Dy_wrt_CO_m, Dpy_wrt_CO_rad, - Dsigma_wrt_CO_m, Ddelta_wrt_CO): + Dzeta_wrt_CO_m, Ddelta_wrt_CO): n_part = 1 for var in [Dx_wrt_CO_m, Dpx_wrt_CO_rad, Dy_wrt_CO_m, Dpy_wrt_CO_rad, - Dsigma_wrt_CO_m, Ddelta_wrt_CO]: + Dzeta_wrt_CO_m, Ddelta_wrt_CO]: if hasattr(var, '__iter__'): if n_part == 1: n_part = len(var) @@ -22,26 +23,26 @@ def vectorize_all_coords(Dx_wrt_CO_m, Dpx_wrt_CO_rad, Dpx_wrt_CO_rad = Dpx_wrt_CO_rad + np.zeros(n_part) Dy_wrt_CO_m = Dy_wrt_CO_m + np.zeros(n_part) Dpy_wrt_CO_rad = Dpy_wrt_CO_rad + np.zeros(n_part) - Dsigma_wrt_CO_m = Dsigma_wrt_CO_m + np.zeros(n_part) + Dzeta_wrt_CO_m = Dzeta_wrt_CO_m + np.zeros(n_part) Ddelta_wrt_CO = Ddelta_wrt_CO + np.zeros(n_part) return Dx_wrt_CO_m, Dpx_wrt_CO_rad,\ Dy_wrt_CO_m, Dpy_wrt_CO_rad,\ - Dsigma_wrt_CO_m, Ddelta_wrt_CO + Dzeta_wrt_CO_m, Ddelta_wrt_CO def track_particle_sixtrack( partCO, Dx_wrt_CO_m, Dpx_wrt_CO_rad, Dy_wrt_CO_m, Dpy_wrt_CO_rad, - Dsigma_wrt_CO_m, Ddelta_wrt_CO, n_turns, + Dzeta_wrt_CO_m, Ddelta_wrt_CO, n_turns, input_folder='./'): Dx_wrt_CO_m, Dpx_wrt_CO_rad,\ Dy_wrt_CO_m, Dpy_wrt_CO_rad,\ - Dsigma_wrt_CO_m, Ddelta_wrt_CO = vectorize_all_coords( + Dzeta_wrt_CO_m, Ddelta_wrt_CO = vectorize_all_coords( Dx_wrt_CO_m, Dpx_wrt_CO_rad, Dy_wrt_CO_m, Dpy_wrt_CO_rad, - Dsigma_wrt_CO_m, Ddelta_wrt_CO) + Dzeta_wrt_CO_m, Ddelta_wrt_CO) n_part = len(Dx_wrt_CO_m) @@ -79,22 +80,22 @@ def track_particle_sixtrack( lines_f13 = [] for i_part in range(n_part): - if type(partCO) is xline.Particles: + if type(partCO) == xp.Particles: temp_part = partCO.copy() else: - temp_part = xline.Particles(**partCO) + temp_part = xp.Particles(**partCO) temp_part.x += Dx_wrt_CO_m[i_part] temp_part.px += Dpx_wrt_CO_rad[i_part] temp_part.y += Dy_wrt_CO_m[i_part] temp_part.py += Dpy_wrt_CO_rad[i_part] - temp_part.sigma += Dsigma_wrt_CO_m[i_part] - temp_part.delta += Ddelta_wrt_CO[i_part] + temp_part.zeta += Dzeta_wrt_CO_m[i_part] + temp_part.update_delta(temp_part.delta + Ddelta_wrt_CO[i_part]) lines_f13.append('%.10e\n' % ((temp_part.x) * 1e3)) lines_f13.append('%.10e\n' % ((temp_part.px) * temp_part.rpp * 1e3)) lines_f13.append('%.10e\n' % ((temp_part.y) * 1e3)) lines_f13.append('%.10e\n' % ((temp_part.py) * temp_part.rpp * 1e3)) - lines_f13.append('%.10e\n' % ((temp_part.sigma) * 1e3)) + lines_f13.append('%.10e\n' % ((temp_part.zeta) * 1e3)) lines_f13.append('%.10e\n' % ((temp_part.delta))) if i_part % 2 == 1: lines_f13.append('%.10e\n' % (temp_part.energy0 * 1e-6)) @@ -144,7 +145,7 @@ def track_particle_sixtrack( px_tbt = np.zeros((n_turns, n_part)) y_tbt = np.zeros((n_turns, n_part)) py_tbt = np.zeros((n_turns, n_part)) - sigma_tbt = np.zeros((n_turns, n_part)) + zeta_tbt = np.zeros((n_turns, n_part)) delta_tbt = np.zeros((n_turns, n_part)) for i_part in range(n_part): @@ -153,86 +154,36 @@ def track_particle_sixtrack( px_tbt[:, i_part] = sixdump_part.px y_tbt[:, i_part] = sixdump_part.y py_tbt[:, i_part] = sixdump_part.py - sigma_tbt[:, i_part] = sixdump_part.sigma + zeta_tbt[:, i_part] = sixdump_part.zeta delta_tbt[:, i_part] = sixdump_part.delta extra = {} - return x_tbt, px_tbt, y_tbt, py_tbt, sigma_tbt, delta_tbt, extra + return x_tbt, px_tbt, y_tbt, py_tbt, zeta_tbt, delta_tbt, extra -def track_particle_xline(line, part, Dx_wrt_CO_m, Dpx_wrt_CO_rad, - Dy_wrt_CO_m, Dpy_wrt_CO_rad, - Dsigma_wrt_CO_m, Ddelta_wrt_CO, n_turns, verbose=False): - - Dx_wrt_CO_m, Dpx_wrt_CO_rad,\ - Dy_wrt_CO_m, Dpy_wrt_CO_rad,\ - Dsigma_wrt_CO_m, Ddelta_wrt_CO = vectorize_all_coords( - Dx_wrt_CO_m, Dpx_wrt_CO_rad, - Dy_wrt_CO_m, Dpy_wrt_CO_rad, - Dsigma_wrt_CO_m, Ddelta_wrt_CO) - - part.x += Dx_wrt_CO_m - part.px += Dpx_wrt_CO_rad - part.y += Dy_wrt_CO_m - part.py += Dpy_wrt_CO_rad - part.sigma += Dsigma_wrt_CO_m - part.delta += Ddelta_wrt_CO - - x_tbt = [] - px_tbt = [] - y_tbt = [] - py_tbt = [] - sigma_tbt = [] - delta_tbt = [] - - for i_turn in range(n_turns): - if verbose: - print('Turn %d/%d' % (i_turn, n_turns)) - - x_tbt.append(part.x.copy()) - px_tbt.append(part.px.copy()) - y_tbt.append(part.y.copy()) - py_tbt.append(part.py.copy()) - sigma_tbt.append(part.sigma.copy()) - delta_tbt.append(part.delta.copy()) - - line.track(part) - - x_tbt = np.array(x_tbt) - px_tbt = np.array(px_tbt) - y_tbt = np.array(y_tbt) - py_tbt = np.array(py_tbt) - sigma_tbt = np.array(sigma_tbt) - delta_tbt = np.array(delta_tbt) - - extra = {} - - return x_tbt, px_tbt, y_tbt, py_tbt, sigma_tbt, delta_tbt def track_particle_xtrack( line, partCO, Dx_wrt_CO_m, Dpx_wrt_CO_rad, Dy_wrt_CO_m, Dpy_wrt_CO_rad, - Dsigma_wrt_CO_m, Ddelta_wrt_CO, n_turns, + Dzeta_wrt_CO_m, Ddelta_wrt_CO, n_turns, _context=None): - Dx_wrt_CO_m, Dpx_wrt_CO_rad,\ Dy_wrt_CO_m, Dpy_wrt_CO_rad,\ - Dsigma_wrt_CO_m, Ddelta_wrt_CO = vectorize_all_coords( + Dzeta_wrt_CO_m, Ddelta_wrt_CO = vectorize_all_coords( Dx_wrt_CO_m, Dpx_wrt_CO_rad, Dy_wrt_CO_m, Dpy_wrt_CO_rad, - Dsigma_wrt_CO_m, Ddelta_wrt_CO) + Dzeta_wrt_CO_m, Ddelta_wrt_CO) - import xtrack as xt - tracker = xt.Tracker(_context=_context, sequence=line) - particles = xt.Particles(_context=_context, + tracker = xt.Tracker(_context=_context, line=line) + particles = xp.Particles(_context=_context, p0c = partCO.p0c, x = partCO.x + Dx_wrt_CO_m, px = partCO.px + Dpx_wrt_CO_rad, y = partCO.y + Dy_wrt_CO_m, py = partCO.py + Dpy_wrt_CO_rad, - zeta = partCO.zeta + Dsigma_wrt_CO_m, + zeta = partCO.zeta + Dzeta_wrt_CO_m, delta = partCO.delta + Ddelta_wrt_CO) print('Start track') @@ -244,14 +195,14 @@ def track_particle_xtrack( px_tbt = tracker.record_last_track.px.copy().T y_tbt = tracker.record_last_track.y.copy().T py_tbt = tracker.record_last_track.py.copy().T - sigma_tbt = tracker.record_last_track.zeta.copy().T + zeta_tbt = tracker.record_last_track.zeta.copy().T delta_tbt = tracker.record_last_track.delta.copy().T print('Done loading!') extra = {'particles': particles, 'tracker': tracker} - return x_tbt, px_tbt, y_tbt, py_tbt, sigma_tbt, delta_tbt, extra + return x_tbt, px_tbt, y_tbt, py_tbt, zeta_tbt, delta_tbt, extra def betafun_from_ellip(x_tbt, px_tbt): diff --git a/python_examples/ions_python/checks_and_doc/t008_check_against_sixtrack.py b/python_examples/ions_python/checks_and_doc/t008_check_against_sixtrack.py index b01da6b..5689b68 100644 --- a/python_examples/ions_python/checks_and_doc/t008_check_against_sixtrack.py +++ b/python_examples/ions_python/checks_and_doc/t008_check_against_sixtrack.py @@ -1,10 +1,11 @@ import json import numpy as np +import NAFFlib import helpers as hp import matplotlib.pyplot as plt -import xline as xl import xtrack as xt +import xpart as xp import sixtracktools @@ -13,40 +14,33 @@ num_turns = 5 -with open('../xlines/line_bb_for_tracking.json', 'r') as fid: +with open('../xsuite_lines/line_bb_for_tracking.json', 'r') as fid: dict_line_xtrack = json.load(fid) -#with open('../xline/line_bb_dipole_cancelled.json', 'r') as fid: -# dict_line_old = json.load(fid) +line = xt.Line.from_dict(dict_line_xtrack) -line = xl.Line.from_dict(dict_line_xtrack) -#line = xl.Line.from_dict(dict_line_old) - -partCO = xl.Particles.from_dict(dict_line_xtrack['particle_on_tracker_co']) +partCO = xp.Particles.from_dict(dict_line_xtrack['particle_on_tracker_co']) (x_tbt_sixtrack, px_tbt_sixtrack, y_tbt_sixtrack, py_tbt_sixtrack, - sigma_tbt_sixtrack, delta_tbt_sixtrack, extra) = hp.track_particle_sixtrack( + zeta_tbt_sixtrack, delta_tbt_sixtrack, extra) = hp.track_particle_sixtrack( partCO=partCO, Dx_wrt_CO_m=np.array(displace_x), Dpx_wrt_CO_rad=0, Dy_wrt_CO_m=np.array(displace_y), Dpy_wrt_CO_rad=0., - Dsigma_wrt_CO_m=0., Ddelta_wrt_CO=0., n_turns=num_turns, + Dzeta_wrt_CO_m=0., Ddelta_wrt_CO=0., n_turns=num_turns, input_folder='../') -tracker = xt.Tracker(sequence=line) - -part_track = partCO.copy() -part_track.x += np.array(displace_x) -part_track.y += np.array(displace_y) +tracker = xt.Tracker(line=line) -particles = xt.Particles(**part_track.to_dict()) -particles.particle_id = np.arange(particles.num_particles) +particles = xp.build_particles(particle_ref=partCO, + x=np.array(displace_x), + y=np.array(displace_y)) tracker.track(particles, turn_by_turn_monitor=True, num_turns=num_turns) print('Xtrack') print(tracker.record_last_track.x) print('Sixtrack') -print(x_tbt_sixtrack.T) +print(x_tbt_sixtrack) assert np.allclose(tracker.record_last_track.x[0, :], x_tbt_sixtrack[:,0], rtol=1e-15, atol=9e-11) From 440f47213f1af6229a90ae13ee92ee06f84ad4d6 Mon Sep 17 00:00:00 2001 From: giadarol Date: Tue, 16 Nov 2021 18:15:26 +0100 Subject: [PATCH 21/22] Cleanup --- .../run3_collisions_python/000_pymask.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/python_examples/run3_collisions_python/000_pymask.py b/python_examples/run3_collisions_python/000_pymask.py index 76dc68c..fa009a2 100644 --- a/python_examples/run3_collisions_python/000_pymask.py +++ b/python_examples/run3_collisions_python/000_pymask.py @@ -194,7 +194,7 @@ print('Leveling not working in this mode!') else: if particle_type == 'ion': # the legacy macro for BB have been checked but not maintained - raise ValueError + raise ValueError # Luminosity levelling vars_for_legacy_level = ['lumi_ip8', 'nco_IP1', 'nco_IP2', 'nco_IP5', 'nco_IP8'] @@ -523,16 +523,15 @@ with open('./optics_orbit_at_start_ring_from_madx.json', 'w') as fid: json.dump(optics_and_co_at_start_ring_from_madx, fid, cls=pm.JEncoder) -################## -# Generate xline # -################## - +######################## +# Generate xtrack line # +######################## if enable_bb_legacy: - print('xline is not generated with bb legacy macros') + print('xtrack line is not generated with bb legacy macros') else: - pm.generate_xline(mad_track, sequence_to_track, bb_df_track, + pm.generate_xsuite_line(mad_track, sequence_to_track, bb_df_track, optics_and_co_at_start_ring_from_madx, - folder_name = './xlines', + folder_name = './xsuite_lines', skip_mad_use=True, prepare_line_for_xtrack=True) From 79362c063005eeec4464bbcb5118da373c90053d Mon Sep 17 00:00:00 2001 From: Giovanni Iadarola Date: Sat, 20 Nov 2021 13:21:40 +0100 Subject: [PATCH 22/22] Added _pkg_root --- pymask/__init__.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/pymask/__init__.py b/pymask/__init__.py index e243297..68aa81d 100644 --- a/pymask/__init__.py +++ b/pymask/__init__.py @@ -7,3 +7,7 @@ from .lumi import * from .coupling import * from .tunechroma import * + +from pathlib import Path +_pkg_root = Path(__file__).parent.absolute() +del(Path)