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) diff --git a/pymask/beambeam.py b/pymask/beambeam.py index 1bbf9d8..efe8024 100644 --- a/pymask/beambeam.py +++ b/pymask/beambeam.py @@ -614,57 +614,67 @@ 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): - 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 + 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.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'] - - 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'] + 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'] + 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 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] + 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.elements[ii] = newee + + def crabbing_strong_beam(mad, bb_dfs, z_crab_twiss, @@ -871,8 +881,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 +889,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, 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 diff --git a/pymask/pymasktools.py b/pymask/pymasktools.py index 36bbc40..0ba19f8 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, @@ -404,95 +402,24 @@ 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_xline(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, - 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 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 +427,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,29 +441,28 @@ 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: - 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'])) + _restore_beam_beam(tracker.line) - particle_on_tracker_co = find_closed_orbit_from_tracker(tracker, - optics_and_co_at_start_ring_from_madx['particle_on_madx_co']) + 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, + _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) - (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) + (WW_finite_diffs, WWInv_finite_diffs, RotMat_finite_diffs + ) = 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 @@ -562,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?!') diff --git a/python_examples/hl_lhc_collisions_python/000_pymask.py b/python_examples/hl_lhc_collisions_python/000_pymask.py index 76dc68c..fa009a2 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'] @@ -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/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/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/t004_compare_xlines.py b/python_examples/hl_lhc_collisions_python/checks_and_doc/t004_check_output_consistency.py similarity index 78% 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 index d6f366e..3d2d8c6 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_check_output_consistency.py @@ -1,10 +1,14 @@ import time import shutil import pickle +import json import numpy as np -import xline import sixtracktools +import xtrack as xt +import xfields as xf + + # Tests b1 with bb @@ -20,9 +24,9 @@ 'strict': False, }, { - 'test_name': 'B1 - pymask xline vs pymask sixtrack input', - 'path_test': '../xlines/line_bb_dipole_not_cancelled.json', - 'type_test': 'xline', + 'test_name': 'B1 - pymask xsuite vs pymask sixtrack input', + 'path_test': '../xsuite_lines/line_bb_dipole_not_cancelled.json', + 'type_test': 'xsuite', 'path_ref': '../', 'type_ref': 'sixtrack', 'rtol': 4e-7, @@ -44,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, @@ -60,17 +64,14 @@ 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) + 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 = xline.Line.from_sixinput(sixinput_test) + ltest = xt.Line.from_sixinput(sixinput_test) print('Done') else: raise ValueError('What?!') @@ -135,8 +136,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 +145,14 @@ def prepare_line(path, input_type): if np.isscalar(dref[kk]) and dtest[kk] == dref[kk]: continue + if isinstance(dref[kk], dict): + 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] val_ref = dref[kk] @@ -172,26 +181,30 @@ 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 # 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 +239,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 +252,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": 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..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 @@ -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' -type_test = 'xline' +#path_test = '../xsuite_lines/line_bb_dipole_not_cancelled.json' +#type_test = 'xsuite' def prepare_line(path, input_type): - if input_type == 'xline': - # Load xline machine - ltest = xline.Line.from_json(path) + 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 = xline.Line.from_sixinput(sixinput_test) + ltest = xt.Line.from_sixinput(sixinput_test) print('Done') else: raise ValueError('What?!') @@ -102,8 +105,8 @@ 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, - 'o', color = 'r', alpha=.5, label='strong pyst') + [getattr(bb, f'delta_{plane}') for bb in bb_elems])+R1_orbit, + '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) @@ -119,12 +122,13 @@ 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]) -for bb in bb_all: bb.enabled = False +bb_all, _ = ltest.get_elements_of_type([xf.BeamBeamBiGaussian2D, + xf.BeamBeamBiGaussian3D]) +for bb in bb_all: + bb.q0 = 0 # # 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 +138,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) @@ -182,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', @@ -220,16 +218,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]) 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..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 @@ -5,11 +5,12 @@ 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' +track_with = 'sixtrack' epsn_x = 2.5e-6 epsn_y = 2.5e-6 @@ -21,24 +22,25 @@ def prepare_line(path, input_type): - if input_type == 'xline': - # Load xline machine - ltest = xline.Line.from_json(path) + 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 = 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='xsuite') -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: @@ -65,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?!') 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..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 @@ -2,31 +2,33 @@ 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], 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 -ampl_sigmas = 1. +ampl_sigmas = 0.2 norm_emit_x = 2.5e-6 geom_emit_x = norm_emit_x / particles.beta0 / particles.gamma0 @@ -35,32 +37,22 @@ 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() +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_test = particles_matched.copy() +tracker.track(particles_test, num_turns=10) -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()) +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) - -tracker.track(particles_matched, num_turns=10) - -plt.plot(particles_matched.x, particles_matched.px) +plt.plot(particles_test.x, particles_test.px) plt.show() 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..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 @@ -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,26 @@ 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) 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) 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) 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)