diff --git a/pymask/__init__.py b/pymask/__init__.py index 4a30a4c..e243297 100644 --- a/pymask/__init__.py +++ b/pymask/__init__.py @@ -1,4 +1,4 @@ -__version__ = "v0.3.1" +__version__ = "v1.1.0" from .beambeam import * from .madpoint import * diff --git a/pymask/beambeam.py b/pymask/beambeam.py index 67fc7b6..1bbf9d8 100644 --- a/pymask/beambeam.py +++ b/pymask/beambeam.py @@ -181,7 +181,8 @@ def generate_set_of_bb_encounters_1beam( harmonic_number = 35640, bunch_spacing_buckets = 10, numberOfHOSlices = 11, - bunch_charge_ppb = 0., + bunch_num_particles = 0., + bunch_particle_charge = 0., sigt=0.0755, relativistic_beta=1., ip_names = ['ip1', 'ip2', 'ip5', 'ip8'], @@ -201,7 +202,8 @@ def generate_set_of_bb_encounters_1beam( if len(myBBLRlist)>0: myBBLR=pd.DataFrame(myBBLRlist)[['beam','other_beam','ip_name','label','identifier']] - myBBLR['self_charge_ppb'] = bunch_charge_ppb + myBBLR['self_num_particles'] = bunch_num_particles + myBBLR['self_particle_charge'] = bunch_particle_charge myBBLR['self_relativistic_beta'] = relativistic_beta myBBLR['elementName']=myBBLR.apply(lambda x: elementName(x.label, x.ip_name.replace('ip', ''), x.beam, x.identifier), axis=1) myBBLR['other_elementName']=myBBLR.apply( @@ -229,7 +231,8 @@ def generate_set_of_bb_encounters_1beam( myBBHO=pd.DataFrame(myBBHOlist)[['beam','other_beam', 'ip_name','label','identifier']] - myBBHO['self_charge_ppb'] = bunch_charge_ppb/numberOfHOSlices + myBBHO['self_num_particles'] = bunch_num_particles/numberOfHOSlices + myBBHO['self_particle_charge'] = bunch_particle_charge myBBHO['self_relativistic_beta'] = relativistic_beta for ip_nn in ip_names: myBBHO.loc[myBBHO['ip_name']==ip_nn, 'atPosition']=list(z_centroids) @@ -249,7 +252,7 @@ def generate_set_of_bb_encounters_1beam( return myBB -def generate_mad_bb_info(bb_df, mode='dummy', madx_reference_bunch_charge=1): +def generate_mad_bb_info(bb_df, mode='dummy', madx_reference_bunch_num_particles=1): if mode == 'dummy': bb_df['elementClass']='beambeam' @@ -259,7 +262,7 @@ def generate_mad_bb_info(bb_df, mode='dummy', madx_reference_bunch_charge=1): 'yma = 1, ' + \ f'charge = 0*{charge}, ' +\ 'slot_id = %d'%({'bb_lr': 4, 'bb_ho': 6}[label]) # need to add 60 for central - bb_df['elementDefinition']=bb_df.apply(lambda x: elementDefinition(x.elementName, x.elementClass, eattributes(x['self_charge_ppb'], x['label'])), axis=1) + bb_df['elementDefinition']=bb_df.apply(lambda x: elementDefinition(x.elementName, x.elementClass, eattributes(x['self_num_particles'], x['label'])), axis=1) bb_df['elementInstallation']=bb_df.apply(lambda x: elementInstallation(x.elementName, x.elementClass, x.atPosition, x.ip_name), axis=1) elif mode=='from_dataframe': bb_df['elementClass']='beambeam' @@ -272,7 +275,7 @@ def generate_mad_bb_info(bb_df, mode='dummy', madx_reference_bunch_charge=1): bb_df['elementDefinition']=bb_df.apply(lambda x: elementDefinition(x.elementName, x.elementClass, eattributes(np.sqrt(x['other_Sigma_11']),np.sqrt(x['other_Sigma_33']), x['xma'], x['yma'], - x['other_charge_ppb']/madx_reference_bunch_charge, x['label'])), + x['other_particle_charge']*x['other_num_particles']/madx_reference_bunch_num_particles, x['label'])), # patch due to the fact that mad-x takes n_part from the weak beam axis=1) bb_df['elementInstallation']=bb_df.apply(lambda x: elementInstallation(x.elementName, x.elementClass, x.atPosition, x.ip_name), axis=1) else: @@ -290,8 +293,10 @@ def get_counter_rotating(bb_df): c_bb_df['identifier'] = bb_df['identifier'] c_bb_df['elementClass'] = bb_df['elementClass'] c_bb_df['elementName'] = bb_df['elementName'] - c_bb_df['self_charge_ppb'] = bb_df['self_charge_ppb'] - c_bb_df['other_charge_ppb'] = bb_df['other_charge_ppb'] + c_bb_df['self_num_particles'] = bb_df['self_num_particles'] + c_bb_df['other_num_particles'] = bb_df['other_num_particles'] + c_bb_df['self_particle_charge'] = bb_df['self_particle_charge'] + c_bb_df['other_particle_charge'] = bb_df['other_particle_charge'] c_bb_df['other_elementName'] = bb_df['other_elementName'] c_bb_df['atPosition'] = bb_df['atPosition'] * (-1.) @@ -349,9 +354,9 @@ def install_lenses_in_sequence(mad, bb_df, sequence_name, regenerate_mad_bb_info_in_df=True): if regenerate_mad_bb_info_in_df: - madx_reference_bunch_charge = mad.sequence[sequence_name].beam.npart + madx_reference_bunch_num_particles = mad.sequence[sequence_name].beam.npart generate_mad_bb_info(bb_df, mode='from_dataframe', - madx_reference_bunch_charge=madx_reference_bunch_charge) + madx_reference_bunch_num_particles=madx_reference_bunch_num_particles) mad.input(bb_df['elementDefinition'].str.cat(sep='\n')) @@ -429,7 +434,8 @@ def get_partner_corrected_position_and_optics(bb_df_b1, bb_df_b2, ip_position_df for ss in _sigma_names: self_df.loc[ee, f'other_Sigma_{ss}'] = other_df.loc[other_ee, f'self_Sigma_{ss}'] # Get charge of other beam - self_df.loc[ee, 'other_charge_ppb'] = other_df.loc[other_ee, 'self_charge_ppb'] + self_df.loc[ee, 'other_num_particles'] = other_df.loc[other_ee, 'self_num_particles'] + self_df.loc[ee, 'other_particle_charge'] = other_df.loc[other_ee, 'self_particle_charge'] self_df.loc[ee, 'other_relativistic_beta'] = other_df.loc[other_ee, 'self_relativistic_beta'] def compute_separations(bb_df): @@ -621,26 +627,27 @@ def setup_beam_beam_in_line( bb_df, bb_coupling=False, ): - import pysixtrack + import xline assert bb_coupling is False # Not implemented for ee, eename in zip(line.elements, line.element_names): - if isinstance(ee, pysixtrack.elements.BeamBeam4D): - ee.charge = bb_df.loc[eename, 'other_charge_ppb'] + 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 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, pysixtrack.elements.BeamBeam6D): + 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_charge_ppb']] + 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'] @@ -735,7 +742,8 @@ def generate_bb_dataframes(mad, harmonic_number=35640, bunch_spacing_buckets=10, numberOfHOSlices=11, - bunch_population_ppb=None, + bunch_num_particles=None, + bunch_particle_charge=None, sigmaz_m=None, z_crab_twiss=0., remove_dummy_lenses=True): @@ -744,13 +752,17 @@ def generate_bb_dataframes(mad, assert mad.sequence.lhcb1.beam[pp] == mad.sequence.lhcb2.beam[pp] circumference = mad.sequence.lhcb1.beam.circ - madx_reference_bunch_charge = mad.sequence.lhcb1.beam.npart + madx_reference_bunch_num_particles = mad.sequence.lhcb1.beam.npart + + if bunch_num_particles is None: + bunch_num_particles = madx_reference_bunch_num_particles + if bunch_particle_charge is None: + bunch_particle_charge = mad.sequence.lhcb1.beam.charge + + + relativistic_gamma = mad.sequence.lhcb1.beam.gamma relativistic_beta = np.sqrt(1 - 1.0 / relativistic_gamma ** 2) - if bunch_population_ppb is not None: - bunch_charge_ppb = bunch_population_ppb - else: - bunch_charge_ppb = madx_reference_bunch_charge if sigmaz_m is not None: sigt = sigmaz_m @@ -760,15 +772,18 @@ def generate_bb_dataframes(mad, bb_df_b1 = generate_set_of_bb_encounters_1beam( circumference, harmonic_number, bunch_spacing_buckets, - numberOfHOSlices, bunch_charge_ppb, sigt, - relativistic_beta, ip_names, numberOfLRPerIRSide, + numberOfHOSlices, + bunch_num_particles, bunch_particle_charge, + sigt, relativistic_beta, ip_names, numberOfLRPerIRSide, beam_name = 'b1', other_beam_name = 'b2') bb_df_b2 = generate_set_of_bb_encounters_1beam( circumference, harmonic_number, bunch_spacing_buckets, - numberOfHOSlices, bunch_charge_ppb, sigt, + numberOfHOSlices, + bunch_num_particles, bunch_particle_charge, + sigt, relativistic_beta, ip_names, numberOfLRPerIRSide, beam_name = 'b2', other_beam_name = 'b1') @@ -808,13 +823,13 @@ def generate_bb_dataframes(mad, # Generate mad info generate_mad_bb_info(bb_df_b1, mode='from_dataframe', - madx_reference_bunch_charge=madx_reference_bunch_charge) + madx_reference_bunch_num_particles=madx_reference_bunch_num_particles) generate_mad_bb_info(bb_df_b2, mode='from_dataframe', - madx_reference_bunch_charge=madx_reference_bunch_charge) + madx_reference_bunch_num_particles=madx_reference_bunch_num_particles) generate_mad_bb_info(bb_df_b3, mode='from_dataframe', - madx_reference_bunch_charge=madx_reference_bunch_charge) + madx_reference_bunch_num_particles=madx_reference_bunch_num_particles) generate_mad_bb_info(bb_df_b4, mode='from_dataframe', - madx_reference_bunch_charge=madx_reference_bunch_charge) + madx_reference_bunch_num_particles=madx_reference_bunch_num_particles) bb_dfs = { 'b1': bb_df_b1, @@ -873,4 +888,4 @@ def compute_xma_yma(bb_df): ) bb_df['xma'] = xma - bb_df['yma'] = yma \ No newline at end of file + bb_df['yma'] = yma diff --git a/pymask/linear_normal_form.py b/pymask/linear_normal_form.py new file mode 100644 index 0000000..5cfdf74 --- /dev/null +++ b/pymask/linear_normal_form.py @@ -0,0 +1,223 @@ +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 2111fa6..36bbc40 100644 --- a/pymask/pymasktools.py +++ b/pymask/pymasktools.py @@ -1,9 +1,25 @@ import os import pickle +import json import numpy as np +import xline as xl +import xtrack as xt + 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): + if isinstance(obj, np.ndarray): + return obj.tolist() + elif np.issubdtype(type(obj), np.integer): + return int(obj) + else: + return json.JSONEncoder.default(self, obj) def make_links(links_dict, force=False): for kk in links_dict.keys(): @@ -182,7 +198,8 @@ def check_separations_against_madvars(checks, twiss_df_b1, twiss_df_b2, variable cc['plane'], target, cc['tol']) def generate_sixtrack_input(mad, seq_name, bb_df, output_folder, - reference_bunch_charge_sixtrack_ppb, + reference_num_particles_sixtrack, + reference_particle_charge_sixtrack, emitnx_sixtrack_um, emitny_sixtrack_um, sigz_sixtrack_m, @@ -193,7 +210,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) @@ -233,8 +250,11 @@ def generate_sixtrack_input(mad, seq_name, bb_df, output_folder, if len(sxt_df_4d) > 0: sxt_df_4d['h-sep [mm]'] = -sxt_df_4d['separation_x']*1e3 sxt_df_4d['v-sep [mm]'] = -sxt_df_4d['separation_y']*1e3 - sxt_df_4d['strength-ratio'] = (sxt_df_4d['other_charge_ppb'] - / reference_bunch_charge_sixtrack_ppb) + sxt_df_4d['strength-ratio'] = ( + sxt_df_4d['other_num_particles'] + * sxt_df_4d['other_particle_charge'] + / reference_num_particles_sixtrack) + #/ reference_particle_charge_sixtrack)# patch for sixtrack inconsistency sxt_df_4d['4dSxx [mm*mm]'] = sxt_df_4d['other_Sigma_11']*1e6 sxt_df_4d['4dSyy [mm*mm]'] = sxt_df_4d['other_Sigma_33']*1e6 sxt_df_4d['4dSxy [mm*mm]'] = sxt_df_4d['other_Sigma_13']*1e6 @@ -256,7 +276,11 @@ def generate_sixtrack_input(mad, seq_name, bb_df, output_folder, sxt_df_6d['v-sep [mm]'] = -sxt_df_6d['separation_y']*1e3 sxt_df_6d['phi [rad]'] = sxt_df_6d['phi'] sxt_df_6d['alpha [rad]'] = sxt_df_6d['alpha'] - sxt_df_6d['strength-ratio'] = sxt_df_6d['other_charge_ppb']/reference_bunch_charge_sixtrack_ppb + sxt_df_6d['strength-ratio'] = ( + sxt_df_6d['other_num_particles'] + * sxt_df_6d['other_particle_charge'] + / reference_num_particles_sixtrack) + #/ reference_particle_charge_sixtrack) # patch for sixtrack sxt_df_6d['Sxx [mm*mm]'] = sxt_df_6d['other_Sigma_11'] *1e6 sxt_df_6d['Sxxp [mm*mrad]'] = sxt_df_6d['other_Sigma_12'] *1e6 sxt_df_6d['Sxpxp [mrad*mrad]'] = sxt_df_6d['other_Sigma_22'] *1e6 @@ -290,7 +314,7 @@ def generate_sixtrack_input(mad, seq_name, bb_df, output_folder, ]), axis=1) f3_common_settings = ' '.join([ - f"{reference_bunch_charge_sixtrack_ppb}", + f"{reference_num_particles_sixtrack*reference_particle_charge_sixtrack}", f"{emitnx_sixtrack_um}", f"{emitny_sixtrack_um}", f"{sigz_sixtrack_m}", @@ -336,19 +360,34 @@ def get_optics_and_orbit_at_start_ring(mad, seq_name, with_bb_forces=False, # Twiss and get closed-orbit if not skip_mad_use: mad.use(sequence=seq_name) - twiss_table = mad.twiss() + twiss_table = mad.twiss(rmatrix=True) if initial_bb_state is not None: mad.globals.on_bb_switch = initial_bb_state - beta0 = mad.sequence[seq_name].beam.beta - gamma0 = mad.sequence[seq_name].beam.gamma - p0c_eV = mad.sequence[seq_name].beam.pc*1.e9 + mad_beam = mad.sequence[seq_name].beam + assert mad_beam.deltap == 0, "Not implemented." + + particle_on_madx_co = xl.Particles( + p0c = mad_beam.pc*1e9, + q0 = mad_beam.charge, + mass0 = mad_beam.mass*1e9, + s = 0, + x = twiss_table.x[0], + px = twiss_table.px[0], + y = twiss_table.y[0], + py = twiss_table.py[0], + tau = twiss_table.t[0], + ptau = twiss_table.pt[0], + ) + + RR_madx = np.zeros([6,6]) - optics_at_start_ring = { - 'beta': beta0, - 'gamma' : gamma0, - 'p0c_eV': p0c_eV, + for ii in range(6): + for jj in range(6): + RR_madx[ii, jj] = getattr(twiss_table, f're{ii+1}{jj+1}')[0] + + optics_and_co_at_start_ring_from_madx = { 'betx': twiss_table.betx[0], 'bety': twiss_table.bety[0], 'alfx': twiss_table.alfx[0], @@ -357,37 +396,103 @@ def get_optics_and_orbit_at_start_ring(mad, seq_name, with_bb_forces=False, 'dy': twiss_table.dy[0], 'dpx': twiss_table.dpx[0], 'dpy': twiss_table.dpy[0], - 'x' : twiss_table.x[0], - 'px' : twiss_table.px[0], - 'y' : twiss_table.y[0], - 'py' : twiss_table.py[0], - 't' : twiss_table.t[0], - 'pt' : twiss_table.pt[0], - #convert tau, pt to sigma,delta - 'sigma' : beta0 * twiss_table.t[0], - 'delta' : ((twiss_table.pt[0]**2 + - 2.*twiss_table.pt[0]/beta0) + 1.)**0.5 - 1. + 'RR_madx': RR_madx, + 'particle_on_madx_co': particle_on_madx_co.to_dict() } - return optics_at_start_ring -def generate_pysixtrack_line_with_bb(mad, seq_name, bb_df, - closed_orbit_method='from_mad', pickle_lines_in_folder=None, - skip_mad_use=False): + return optics_and_co_at_start_ring_from_madx + + + + +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) - opt_and_CO = get_optics_and_orbit_at_start_ring(mad, seq_name, - with_bb_forces=False, skip_mad_use=True) - # Build pysixtrack model - import pysixtrack - pysxt_line = pysixtrack.Line.from_madx_sequence( - mad.sequence[seq_name]) +def generate_xline(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( + mad.sequence[seq_name], apply_madx_errors=True) + print('Done building xline.') if bb_df is not None: - bb.setup_beam_beam_in_line(pysxt_line, bb_df, bb_coupling=False) + bb.setup_beam_beam_in_line(line, bb_df, bb_coupling=False) # Temporary fix due to bug in mad loader - cavities, cav_names = pysxt_line.get_elements_of_type( - pysixtrack.elements.Cavity) + cavities, cav_names = line.get_elements_of_type( + xl.elements.Cavity) for cc, nn in zip(cavities, cav_names): if cc.frequency ==0.: ii_mad = mad.sequence[seq_name].element_names().index(nn) @@ -395,39 +500,65 @@ def generate_pysixtrack_line_with_bb(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 - mad_CO = np.array([opt_and_CO[kk] for kk in ['x', 'px', 'y', 'py', 'sigma', 'delta']]) - - pysxt_line.disable_beambeam() - part_on_CO = pysxt_line.find_closed_orbit( - guess=mad_CO, p0c=opt_and_CO['p0c_eV'], - method={'from_mad': 'get_guess', 'from_tracking': 'Nelder-Mead'}[closed_orbit_method]) - pysxt_line.enable_beambeam() - - pysxt_line_bb_dipole_cancelled = pysxt_line.copy() - - pysxt_line_bb_dipole_cancelled.beambeam_store_closed_orbit_and_dipolar_kicks( - part_on_CO, - separation_given_wrt_closed_orbit_4D=True, - separation_given_wrt_closed_orbit_6D=True) - - pysxt_dict = { - 'line_bb_dipole_not_cancelled': pysxt_line, - 'line_bb_dipole_cancelled': pysxt_line_bb_dipole_cancelled, - 'particle_on_closed_orbit': part_on_CO} - - if pickle_lines_in_folder is not None: - pysix_fol_name = pickle_lines_in_folder - os.makedirs(pysix_fol_name, exist_ok=True) - - with open(pysix_fol_name + "/line_bb_dipole_not_cancelled.pkl", "wb") as fid: - pickle.dump(pysxt_line.to_dict(keepextra=True), fid) - - with open(pysix_fol_name + "/line_bb_dipole_cancelled.pkl", "wb") as fid: - pickle.dump(pysxt_line_bb_dipole_cancelled.to_dict(keepextra=True), fid) - - with open(pysix_fol_name + "/particle_on_closed_orbit.pkl", "wb") as fid: - pickle.dump(part_on_CO.to_dict(), fid) - - return pysxt_dict - - + line_bb_dipole_not_cancelled_dict = line.to_dict(keepextra=True) + 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'] = ( + optics_and_co_at_start_ring_from_madx['RR_madx']) + + if folder_name is not None: + os.makedirs(folder_name, exist_ok=True) + # Note that full separation and not strong beam position is present + # in bb lenses (for comparison with sixtrack input) + with open(folder_name + '/line_bb_dipole_not_cancelled.json', 'w') as fid: + json.dump(line_bb_dipole_not_cancelled_dict, fid, cls=JEncoder) + + if prepare_line_for_xtrack: + tracker = xt.Tracker(sequence=line) + + # Disable beam-beam + for ee in tracker.line.elements: + if ee.__class__.__name__.startswith('BeamBeam'): + 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']) + + RR_finite_diffs = compute_R_matrix_finite_differences( + particle_on_tracker_co, tracker, symplectify=True, + **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) + + line_bb_for_tracking_dict = line.to_dict(keepextra=True) + 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 + line_bb_for_tracking_dict['WW_finite_diffs'] = WW_finite_diffs + line_bb_for_tracking_dict['WWInv_finite_diffs'] = WWInv_finite_diffs + line_bb_for_tracking_dict['RotMat_finite_diffs'] = RotMat_finite_diffs + + if folder_name is not None: + os.makedirs(folder_name, exist_ok=True) + with open(folder_name + + '/line_bb_for_tracking.json', 'w') as fid: + json.dump(line_bb_for_tracking_dict, fid, cls=JEncoder) + + +def save_mad_sequence_and_error(mad, seq_name, filename='lhc'): + mad.select(flag="error",clear=True) + mad.select(flag="error",class_="multipole") + mad.select(flag="error",class_="hkicker") + mad.select(flag="error",class_="vkicker") + mad.select(flag="error",class_="kicker") + mad.esave(file=filename + "_errors.tfs") + mad.select(flag="error",clear=True) + mad.select(flag="error",full=True) + mad.esave(file=filename + "_errors_all.tfs") + mad.save(sequence=seq_name,beam=True,file=filename + "_seq.madx") diff --git a/python_examples/clean_all b/python_examples/clean_all index 996393d..bf1b8aa 100644 --- a/python_examples/clean_all +++ b/python_examples/clean_all @@ -1,15 +1,18 @@ rm fc.* rm *.log +rm final_seq.madx +rm *.tfs rm bb_lenses.dat rm *.parquet rm *.pkl +rm *.json rm twiss_*.b1 rm twiss_*.b2 rm twiss_bb rm last_twiss.*.gz rm fort.* rm *.tfs -rm -r pysixtrack +rm -r xline rm -r temp rm -r __pycache__ diff --git a/python_examples/hl_lhc_collisions_python/000_pymask.py b/python_examples/hl_lhc_collisions_python/000_pymask.py index 0f37d14..76dc68c 100644 --- a/python_examples/hl_lhc_collisions_python/000_pymask.py +++ b/python_examples/hl_lhc_collisions_python/000_pymask.py @@ -1,6 +1,6 @@ import os import sys -import pickle +import json import numpy as np @@ -89,19 +89,19 @@ mad.globals.nrj = configuration['beam_energy_tot'] particle_type = 'proton' -if 'beam_mass' in configuration.keys(): - beam_mass = configuration['beam_mass'] +if 'particle_mass' in configuration.keys(): + particle_mass = configuration['particle_mass'] particle_type = 'ion' else: - beam_mass = mad.globals.pmass # proton mass + particle_mass = mad.globals.pmass # proton mass -if 'beam_charge' in configuration.keys(): - beam_charge = configuration['beam_charge'] +if 'particle_charge' in configuration.keys(): + particle_charge = configuration['particle_charge'] particle_type = 'ion' else: - beam_charge = 1. + particle_charge = 1. -gamma_rel = (beam_charge*configuration['beam_energy_tot'])/beam_mass +gamma_rel = (particle_charge*configuration['beam_energy_tot'])/particle_mass for ss in mad.sequence.keys(): # bv and bv_aux flags if ss == 'lhcb1': @@ -115,15 +115,15 @@ mad.globals['bv_aux'] = ss_bv_aux mad.input(f''' beam, particle={particle_type},sequence={ss}, - energy={configuration['beam_energy_tot']*beam_charge}, + energy={configuration['beam_energy_tot']*particle_charge}, sigt={configuration['beam_sigt']}, bv={ss_beam_bv}, npart={configuration['beam_npart']}, sige={configuration['beam_sige']}, ex={configuration['beam_norm_emit_x'] * 1e-6 / gamma_rel}, ey={configuration['beam_norm_emit_y'] * 1e-6 / gamma_rel}, - mass={beam_mass}, - charge={beam_charge}, + mass={particle_mass}, + charge={particle_charge}, ''') @@ -247,20 +247,21 @@ # Prepare bb dataframes if enable_bb_python: + bbconfig = configuration['beambeam_config'] bb_dfs = pm.generate_bb_dataframes(mad, ip_names=['ip1', 'ip2', 'ip5', 'ip8'], harmonic_number=35640, - numberOfLRPerIRSide=configuration['numberOfLRPerIRSide'], - bunch_spacing_buckets=configuration['bunch_spacing_buckets'], - numberOfHOSlices=configuration['numberOfHOSlices'], - bunch_population_ppb=configuration['bunch_population_ppb'], - sigmaz_m=configuration['sigmaz_m'], - z_crab_twiss=configuration['z_crab_twiss']*float(enable_crabs), + numberOfLRPerIRSide=bbconfig['numberOfLRPerIRSide'], + bunch_spacing_buckets=bbconfig['bunch_spacing_buckets'], + numberOfHOSlices=bbconfig['numberOfHOSlices'], + bunch_num_particles = bbconfig['bunch_num_particles'], + bunch_particle_charge = bbconfig['bunch_particle_charge'], + sigmaz_m=bbconfig['sigmaz_m'], + z_crab_twiss=bbconfig['z_crab_twiss']*float(enable_crabs), remove_dummy_lenses=True) # Here the datafremes can be edited, e.g. to set bbb intensity - ################### # Generate beam 4 # ################### @@ -286,11 +287,12 @@ check_betas_at_ips=check_betas_at_ips, check_separations_at_ips=False) # For B1, to be generalized for B4 -if 'filling_scheme_json' in configuration.keys(): - filling_scheme_json = configuration['filling_scheme_json'] - bunch_to_track = configuration['bunch_to_track'] +if 'filling_scheme_json' in configuration['beambeam_config'].keys(): + assert 'b4' not in mode + filling_scheme_json = configuration['beambeam_config']['filling_scheme_json'] + bunch_to_track = configuration['beambeam_config']['bunch_to_track'] bb_schedule_to_track_b1 = ost.create_bb_shedule_to_track( - filling_scheme_json,bunch_to_track, beam='1') + filling_scheme_json,bunch_to_track, beam=1) bb_dfs['b1']=ost.filter_bb_df(bb_dfs['b1'],bb_schedule_to_track_b1) ################################################## @@ -337,6 +339,7 @@ # Legacy bb macros if enable_bb_legacy: + bbconfig = configuration['beambeam_config'] assert(beam_to_configure == 1) assert(not(track_from_b4_mad_instance)) assert(not(enable_bb_python)) @@ -344,7 +347,7 @@ mad_track.set_variables_from_dict( params=configuration['pars_for_legacy_bb_macros']) mad_track.set_variables_from_dict( - params={f'par_nho_ir{ir}':configuration['numberOfHOSlices'] + params={f'par_nho_ir{ir}': bbconfig['numberOfHOSlices'] for ir in [1,2,5,8]}) mad_track.input("call, file='modules/module_03_beambeam.madx';") @@ -492,8 +495,9 @@ seq_name=sequence_to_track, bb_df=bb_df_track, output_folder='./', - reference_bunch_charge_sixtrack_ppb=( + reference_num_particles_sixtrack=( mad_track.sequence[sequence_to_track].beam.npart), + reference_particle_charge_sixtrack=mad_track.sequence[sequence_to_track].beam.charge, emitnx_sixtrack_um=( mad_track.sequence[sequence_to_track].beam.exn), emitny_sixtrack_um=( @@ -514,29 +518,27 @@ # Save optics and orbit at start ring # ####################################### -optics_orbit_start_ring = pm.get_optics_and_orbit_at_start_ring( +optics_and_co_at_start_ring_from_madx = pm.get_optics_and_orbit_at_start_ring( mad_track, sequence_to_track, skip_mad_use=True) -with open('./optics_orbit_at_start_ring.pkl', 'wb') as fid: - pickle.dump(optics_orbit_start_ring, fid) - +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 pysixtrack lines # -############################# +################## +# Generate xline # +################## if enable_bb_legacy: - print('Pysixtrack line is not generated with bb legacy macros') + print('xline is not generated with bb legacy macros') else: - pysix_fol_name = "./pysixtrack" - dct_pysxt = pm.generate_pysixtrack_line_with_bb(mad_track, - sequence_to_track, bb_df_track, - closed_orbit_method='from_mad', - pickle_lines_in_folder=pysix_fol_name, - skip_mad_use=True) + pm.generate_xline(mad_track, sequence_to_track, bb_df_track, + optics_and_co_at_start_ring_from_madx, + folder_name = './xlines', + skip_mad_use=True, + prepare_line_for_xtrack=True) -############################# -# Save final twiss # -############################# +################################### +# Save final twiss # +################################### mad_track.globals.on_bb_charge = 0 mad_track.twiss() @@ -544,3 +546,17 @@ sdf = mad_track.get_summ_df('summ') tdf.to_parquet('final_twiss_BBOFF.parquet') sdf.to_parquet('final_summ_BBOFF.parquet') + + +mad_track.globals.on_bb_charge = 1 +mad_track.twiss() +tdf = mad_track.get_twiss_df('twiss') +sdf = mad_track.get_summ_df('summ') +tdf.to_parquet('final_twiss_BBON.parquet') +sdf.to_parquet('final_summ_BBON.parquet') + +############################# +# Save sequence and errors # +############################# +# N.B. this erases the errors in the mad_track instance +# pm.save_mad_sequence_and_error(mad_track, sequence_to_track, filename='final') diff --git a/python_examples/hl_lhc_collisions_python/001_reload.py b/python_examples/hl_lhc_collisions_python/001_reload.py new file mode 100644 index 0000000..b5c7f6a --- /dev/null +++ b/python_examples/hl_lhc_collisions_python/001_reload.py @@ -0,0 +1,18 @@ +import sys +import numpy as np +import xline as xl +import xtrack as xt + +sys.path.append('./modules') +import pymask as pm + +Madx = pm.Madxp +mad = Madx(command_log="mad_final.log") +mad.call("final_seq.madx") +mad.use(sequence="lhcb1") +mad.twiss() +mad.readtable(file="final_errors.tfs", table="errtab") +mad.seterr(table="errtab") +mad.set(format=".15g") +mad.twiss(rmatrix = True) + 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 2d67edb..8caf3b0 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,7 +1,7 @@ import numpy as np import os import sixtracktools -import pysixtrack +import xline import time @@ -79,10 +79,10 @@ def track_particle_sixtrack( lines_f13 = [] for i_part in range(n_part): - if type(partCO) is pysixtrack.Particles: + if type(partCO) is xline.Particles: temp_part = partCO.copy() else: - temp_part = pysixtrack.Particles(**partCO) + temp_part = xline.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] @@ -156,10 +156,12 @@ def track_particle_sixtrack( sigma_tbt[:, i_part] = sixdump_part.sigma delta_tbt[:, i_part] = sixdump_part.delta - return x_tbt, px_tbt, y_tbt, py_tbt, sigma_tbt, delta_tbt + extra = {} + + return x_tbt, px_tbt, y_tbt, py_tbt, sigma_tbt, delta_tbt, extra -def track_particle_pysixtrack(line, part, Dx_wrt_CO_m, Dpx_wrt_CO_rad, +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): @@ -204,13 +206,15 @@ def track_particle_pysixtrack(line, part, Dx_wrt_CO_m, Dpx_wrt_CO_rad, 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_sixtracklib( +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, - device=None): + _context=None): Dx_wrt_CO_m, Dpx_wrt_CO_rad,\ @@ -220,348 +224,35 @@ def track_particle_sixtracklib( Dy_wrt_CO_m, Dpy_wrt_CO_rad, Dsigma_wrt_CO_m, Ddelta_wrt_CO) - #if type(partCO) is pysixtrack.Particles: - # temp_part = partCO.copy() - #else: - # temp_part = pysixtrack.Particles(**partCO) - - import sixtracklib - elements=sixtracklib.Elements() - elements.BeamMonitor(num_stores=n_turns) - elements.append_line(line) - - n_part = len(Dx_wrt_CO_m) - - # Build PyST particle - - ps = sixtracklib.ParticlesSet() - p = ps.Particles(num_particles=n_part) - - for i_part in range(n_part): - - if type(partCO) is pysixtrack.Particles: - part = partCO.copy() - else: - part = pysixtrack.Particles(**partCO) - part.x += Dx_wrt_CO_m[i_part] - part.px += Dpx_wrt_CO_rad[i_part] - part.y += Dy_wrt_CO_m[i_part] - part.py += Dpy_wrt_CO_rad[i_part] - part.sigma += Dsigma_wrt_CO_m[i_part] - part.delta += Ddelta_wrt_CO[i_part] - - part.partid = i_part - part.state = 1 - part.elemid = 0 - part.turn = 0 - - p.from_pysixtrack(part, i_part) - - if device is None: - job = sixtracklib.TrackJob(elements, ps) - else: - job = sixtracklib.TrackJob(elements, ps, device=device) - - job.track(n_turns) - job.collect() - - res = job.output - - #print(res.particles[0]) - x_tbt = res.particles[0].x.reshape(n_turns, n_part) - px_tbt = res.particles[0].px.reshape(n_turns, n_part) - y_tbt = res.particles[0].y.reshape(n_turns, n_part) - py_tbt = res.particles[0].py.reshape(n_turns, n_part) - sigma_tbt = res.particles[0].sigma.reshape(n_turns, n_part) - delta_tbt = res.particles[0].delta.reshape(n_turns, n_part) - - # For now data are saved at the end of the turn by STlib and at the beginning by the others - #x_tbt[1:, :] = x_tbt[:-1, :] - #px_tbt[1:, :] = px_tbt[:-1, :] - #y_tbt[1:, :] = y_tbt[:-1, :] - #py_tbt[1:, :] = py_tbt[:-1, :] - #sigma_tbt[1:, :] = sigma_tbt[:-1, :] - #delta_tbt[1:, :] = delta_tbt[:-1, :] - #x_tbt[0, :] = p.x - #px_tbt[0, :] = p.px - #y_tbt[0, :] = p.y - #py_tbt[0, :] = p.py - #sigma_tbt[0, :] = p.sigma - #delta_tbt[0, :] = p.delta - - print('Done loading!') - - return x_tbt, px_tbt, y_tbt, py_tbt, sigma_tbt, delta_tbt - - -def track_particle_sixtracklib_long( - line, partCO, Dx_wrt_CO_m, Dpx_wrt_CO_rad, - Dy_wrt_CO_m, Dpy_wrt_CO_rad, - Dzeta_wrt_CO_m, Ddelta_wrt_CO, n_turns, - device=None): - - Dx_wrt_CO_m, Dpx_wrt_CO_rad,\ - Dy_wrt_CO_m, Dpy_wrt_CO_rad,\ - 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, - Dzeta_wrt_CO_m, Ddelta_wrt_CO) - - - #if type(partCO) is pysixtrack.Particles: - # part = partCO.copy() - #else: - # part = pysixtrack.Particles(**partCO) - n_turns_tbt=1000 - n_turns_to_store=1000 - skip_turns=int(n_turns)//n_turns_to_store - if skip_turns == 0: - skip_turns = 1 - - - import sixtracklib - elements=sixtracklib.Elements() - #sixtracklib.append_beam_monitors_to_lattice(beam_elements_buffer=elements.cbuffer, - # until_turn_elem_by_elem=0, - # until_turn_turn_by_turn=n_turns_tbt, - # until_turn=n_turns, - # skip_turns=skip_turns - # ) - elements.BeamMonitor(num_stores=n_turns_tbt,start=0,skip=1,is_rolling=False) - elements.BeamMonitor(num_stores=n_turns_to_store,start=0,skip=skip_turns,is_rolling=False) - elements.BeamMonitor(num_stores=1,start=0,skip=1,is_rolling=True) - print(elements.get_elements()) - #elements.BeamMonitor(num_stores=n_turns) - #elements.BeamMonitor(num_stores=n_turns_to_store) - elements.append_line(line) - - n_stores0=elements.get_elements()[0].num_stores - n_stores1=elements.get_elements()[1].num_stores - n_stores2=elements.get_elements()[2].num_stores - n_part = len(Dx_wrt_CO_m) - - # Build PyST particle - - ps = sixtracklib.ParticlesSet() - p = ps.Particles(num_particles=n_part) - - for i_part in range(n_part): - - if type(partCO) is pysixtrack.Particles: - part = partCO.copy() - else: - part = pysixtrack.Particles(**partCO) - part.x += Dx_wrt_CO_m[i_part] - part.px += Dpx_wrt_CO_rad[i_part] - part.y += Dy_wrt_CO_m[i_part] - part.py += Dpy_wrt_CO_rad[i_part] - part.zeta += Dzeta_wrt_CO_m[i_part] - part.delta += Ddelta_wrt_CO[i_part] - - part.partid = i_part - part.state = 1 - part.elemid = 0 - part.turn = 0 - - p.from_pysixtrack(part, i_part) - - if device is None: - job = sixtracklib.TrackJob(elements, ps) - else: - job = sixtracklib.TrackJob(elements, ps, device=device) - - start_tracking_time = time.time() - job.track(n_turns) - end_tracking_time = time.time() - job.collect() - end_collecting_time = time.time() - res = job.output + import xtrack as xt + tracker = xt.Tracker(_context=_context, sequence=line) + particles = xt.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, + delta = partCO.delta + Ddelta_wrt_CO) + + print('Start track') + tracker.track(particles, num_turns=n_turns, turn_by_turn_monitor=True) + print('Done track') #print(res.particles[0]) - #print(res.particles[1]) - - x_tbt_first = res.particles[0].x.reshape(n_stores0,n_part) - px_tbt_first = res.particles[0].px.reshape(n_stores0,n_part) - y_tbt_first = res.particles[0].y.reshape(n_stores0,n_part) - py_tbt_first = res.particles[0].py.reshape(n_stores0,n_part) - zeta_tbt_first = res.particles[0].zeta.reshape(n_stores0,n_part) - delta_tbt_first = res.particles[0].delta.reshape(n_stores0,n_part) - at_turn_tbt_first = res.particles[0].at_turn.reshape(n_stores0,n_part) - state_tbt_first = res.particles[0].state.reshape(n_stores0,n_part) - - x_skip = res.particles[1].x.reshape(n_stores1,n_part) - px_skip = res.particles[1].px.reshape(n_stores1,n_part) - y_skip = res.particles[1].y.reshape(n_stores1,n_part) - py_skip = res.particles[1].py.reshape(n_stores1,n_part) - zeta_skip = res.particles[1].zeta.reshape(n_stores1,n_part) - delta_skip = res.particles[1].delta.reshape(n_stores1,n_part) - at_turn_skip = res.particles[1].at_turn.reshape(n_stores1,n_part) - state_skip = res.particles[1].state.reshape(n_stores1,n_part) - - x_last = res.particles[2].x.reshape(n_stores2,n_part) - px_last = res.particles[2].px.reshape(n_stores2,n_part) - y_last = res.particles[2].y.reshape(n_stores2,n_part) - py_last = res.particles[2].py.reshape(n_stores2,n_part) - zeta_last = res.particles[2].zeta.reshape(n_stores2,n_part) - delta_last = res.particles[2].delta.reshape(n_stores2,n_part) - at_turn_last = res.particles[2].at_turn.reshape(n_stores2,n_part) - state_last = res.particles[2].state.reshape(n_stores2,n_part) - - output_dict = {'x_tbt_first' : x_tbt_first, - 'px_tbt_first' : px_tbt_first, - 'y_tbt_first' : y_tbt_first, - 'py_tbt_first' : py_tbt_first, - 'zeta_tbt_first' : zeta_tbt_first, - 'delta_tbt_first' : delta_tbt_first, - 'at_turn_tbt_first' : at_turn_tbt_first, -# 'state_tbt_first' : state_tbt_first, - 'x_skip' : x_skip, - 'px_skip' : px_skip, - 'y_skip' : y_skip, - 'py_skip' : py_skip, - 'zeta_skip' : zeta_skip, - 'delta_skip' : delta_skip, - 'at_turn_skip': at_turn_skip, - 'state_skip' : state_skip, - 'x_last' : x_last, - 'px_last' : px_last, - 'y_last' : y_last, - 'py_last' : py_last, - 'zeta_last' : zeta_last, - 'delta_last' : delta_last, - 'at_turn_last' : at_turn_last, -# 'state_tbt_last' : state_tbt_last, - 'tracking_time_mins' : (end_tracking_time - start_tracking_time)/60., - 'collecting_time_mins' : (end_collecting_time - end_tracking_time)/60., - } - + x_tbt = tracker.record_last_track.x.copy().T + 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 + delta_tbt = tracker.record_last_track.delta.copy().T print('Done loading!') - return output_dict - -def track_particle_sixtracklib_firstlast( - 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, - device=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( - Dx_wrt_CO_m, Dpx_wrt_CO_rad, - Dy_wrt_CO_m, Dpy_wrt_CO_rad, - Dsigma_wrt_CO_m, Ddelta_wrt_CO) + extra = {'particles': particles, 'tracker': tracker} - #if type(partCO) is pysixtrack.Particles: - # part = partCO.copy() - #else: - # part = pysixtrack.Particles(**partCO) - - n_turns_to_store=1000 - n_turns_tbt=1000 - #skip_turns=1000 - - - import sixtracklib - elements=sixtracklib.Elements() - #sixtracklib.append_beam_monitors_to_lattice(beam_elements_buffer=elements.cbuffer, - # until_turn_elem_by_elem=0, - # until_turn_turn_by_turn=n_turns_tbt, - # until_turn=n_turns, - # skip_turns=skip_turns - # ) - elements.BeamMonitor(num_stores=n_turns_tbt,start=0,skip=1,is_rolling=False) - elements.BeamMonitor(num_stores=n_turns_to_store,start=0,skip=1,is_rolling=True) - print(elements.get_elements()) - #elements.BeamMonitor(num_stores=n_turns) - #elements.BeamMonitor(num_stores=n_turns_to_store) - elements.append_line(line) - - n_stores=elements.get_elements()[1].num_stores - n_part = len(Dx_wrt_CO_m) - - # Build PyST particle + return x_tbt, px_tbt, y_tbt, py_tbt, sigma_tbt, delta_tbt, extra - ps = sixtracklib.ParticlesSet() - p = ps.Particles(num_particles=n_part) - - for i_part in range(n_part): - - if type(partCO) is pysixtrack.Particles: - part = partCO.copy() - else: - part = pysixtrack.Particles(**partCO) - part.x += Dx_wrt_CO_m[i_part] - part.px += Dpx_wrt_CO_rad[i_part] - part.y += Dy_wrt_CO_m[i_part] - part.py += Dpy_wrt_CO_rad[i_part] - part.sigma += Dsigma_wrt_CO_m[i_part] - part.delta += Ddelta_wrt_CO[i_part] - - part.partid = i_part - part.state = 1 - part.elemid = 0 - part.turn = 0 - - p.from_pysixtrack(part, i_part) - - if device is None: - job = sixtracklib.TrackJob(elements, ps) - else: - job = sixtracklib.TrackJob(elements, ps, device=device) - - start_tracking_time = time.time() - job.track(n_turns) - end_tracking_time = time.time() - job.collect() - end_collecting_time = time.time() - res = job.output - - print(res.particles[0]) - print(res.particles[1]) - - x_tbt_first = res.particles[0].x.reshape(n_turns_tbt,n_part) - px_tbt_first = res.particles[0].px.reshape(n_turns_tbt,n_part) - y_tbt_first = res.particles[0].y.reshape(n_turns_tbt,n_part) - py_tbt_first = res.particles[0].py.reshape(n_turns_tbt,n_part) - zeta_tbt_first = res.particles[0].zeta.reshape(n_turns_tbt,n_part) - delta_tbt_first = res.particles[0].delta.reshape(n_turns_tbt,n_part) - at_turn_tbt_first = res.particles[0].at_turn.reshape(n_turns_tbt,n_part) - state_tbt_first = res.particles[0].state.reshape(n_turns_tbt,n_part) - - x_tbt_last = res.particles[1].x.reshape(n_stores,n_part) - px_tbt_last = res.particles[1].px.reshape(n_stores,n_part) - y_tbt_last = res.particles[1].y.reshape(n_stores,n_part) - py_tbt_last = res.particles[1].py.reshape(n_stores,n_part) - zeta_tbt_last = res.particles[1].zeta.reshape(n_stores,n_part) - delta_tbt_last = res.particles[1].delta.reshape(n_stores,n_part) - at_turn_tbt_last = res.particles[1].at_turn.reshape(n_stores,n_part) - state_tbt_last = res.particles[1].state.reshape(n_stores,n_part) - - output_dict = {'x_tbt_first' : x_tbt_first, - 'px_tbt_first' : px_tbt_first, - 'y_tbt_first' : y_tbt_first, - 'py_tbt_first' : py_tbt_first, - 'zeta_tbt_first' : zeta_tbt_first, - 'delta_tbt_first' : delta_tbt_first, - 'at_turn_tbt_first' : at_turn_tbt_first, -# 'state_tbt_first' : state_tbt_first, - 'x_tbt_last' : x_tbt_last, - 'px_tbt_last' : px_tbt_last, - 'y_tbt_last' : y_tbt_last, - 'py_tbt_last' : py_tbt_last, - 'zeta_tbt_last' : zeta_tbt_last, - 'delta_tbt_last' : delta_tbt_last, - 'at_turn_tbt_last' : at_turn_tbt_last, -# 'state_tbt_last' : state_tbt_last, - 'tracking_time_mins' : (end_tracking_time - start_tracking_time)/60., - 'collecting_time_mins' : (end_collecting_time - end_tracking_time)/60., - } - - - print('Done loading!') - return output_dict def betafun_from_ellip(x_tbt, px_tbt): diff --git a/python_examples/hl_lhc_collisions_python/checks_and_doc/t004_compare_pysixtrack_lines.py b/python_examples/hl_lhc_collisions_python/checks_and_doc/t004_compare_pysixtrack_lines.py deleted file mode 100644 index 61c0466..0000000 --- a/python_examples/hl_lhc_collisions_python/checks_and_doc/t004_compare_pysixtrack_lines.py +++ /dev/null @@ -1,162 +0,0 @@ -import shutil -import pickle -import numpy as np - -import pysixtrack -import sixtracktools - -# Test b1 -path_test = '../' -type_test = 'sixtrack' -path_ref = '../../../examples/hl_lhc_collision' -type_ref = 'sixtrack' - -# # Test b4 nobb sixtrack -# path_test = '../../' -# type_test = 'sixtrack' -# path_ref = '../../../examples/hl_lhc_collision_nobb_b4' -# type_ref = 'sixtrack' - -# # Test b4 nobb pysixtrack (not working for now) -# path_test = './pysixtrack/line_bb_dipole_not_cancelled.pkl' -# type_test = 'pysixtrack' -# path_ref = '../hl_lhc_collision_nobb_b4' -# type_ref = 'sixtrack' - -def prepare_line(path, input_type): - - if input_type == 'pysixtrack': - # Load pysixtrack machine - with open(path, "rb") as fid: - ltest = pysixtrack.Line.from_dict(pickle.load(fid)) - elif input_type == 'sixtrack': - print('Build pysixtrack from sixtrack input:') - sixinput_test = sixtracktools.sixinput.SixInput(path) - ltest = pysixtrack.Line.from_sixinput(sixinput_test) - print('Done') - else: - raise ValueError('What?!') - - return ltest - -# Load -ltest = prepare_line(path_test, type_test) -lref = prepare_line(path_ref, type_ref) - -original_length = ltest.get_length() -assert (lref.get_length() - original_length) < 1e-6 - -# Simplify the two machines -for ll in (ltest, lref): - ll.remove_inactive_multipoles(inplace=True) - ll.remove_zero_length_drifts(inplace=True) - ll.merge_consecutive_drifts(inplace=True) - -# Check that the two machines are identical -assert len(ltest) == len(lref) - -assert (ltest.get_length() - original_length) < 1e-6 -assert (lref.get_length() - original_length) < 1e-6 - - -def norm(x): - return np.sqrt(np.sum(np.array(x) ** 2)) - - -for ii, (ee_test, ee_six, nn_test, nn_six) in enumerate( - zip(ltest.elements, lref.elements, ltest.element_names, lref.element_names) -): - assert type(ee_test) == type(ee_six) - - - dtest = ee_test.to_dict(keepextra=True) - dsix = ee_six.to_dict(keepextra=True) - - for kk in dtest.keys(): - - # Check if they are identical - if dtest[kk] == dsix[kk]: - continue - - # Check if the relative error is small - try: - diff_rel = norm(np.array(dtest[kk]) - np.array(dsix[kk])) / norm( - dtest[kk] - ) - except ZeroDivisionError: - diff_rel = 100.0 - if diff_rel < 3e-5: - continue - - # Check if absolute error is small - diff_abs = norm(np.array(dtest[kk]) - np.array(dsix[kk])) - if diff_abs > 0: - print(f"{nn_test}[{kk}] - test:{dtest[kk]} six:{dsix[kk]}") - if diff_abs < 1e-12: - continue - - # Exception: correctors involved in lumi leveling - if nn_test in [ - 'mcbcv.5l8.b1', 'mcbyv.a4l8.b1', 'mcbxv.3l8', - 'mcbyv.4r8.b1', 'mcbyv.b5r8.b1', - 'mcbyh.b5l2.b1', 'mcbyh.4l2.b1', 'mcbxh.3l2', 'mcbyh.a4r2.b1', - 'mcbch.5r2.b1', - 'mcbcv.5l8.b2', 'mcbyv.a4l8.b2', 'mcbxv.3l8', - 'mcbyv.4r8.b2', 'mcbyv.b5r8.b2', - 'mcbyh.b5l2.b2', 'mcbyh.4l2.b2', 'mcbxh.3l2', 'mcbyh.a4r2.b2', - 'mcbch.5r2.b2', 'mcbch.a5r2.b2', 'mcbyh.4r2.b2', 'mcbxh.3r2', - 'mcbyh.a4l2.b2', 'mcbyh.5l2.b2', 'mcbyv.5r8.b2', 'mcbyv.a4r8.b2', - 'mcbxv.3r8', 'mcbyv.4l8.b2', 'mcbcv.b5l8.b2']: - if diff_rel < 1e-2: - continue - - # Exception: drift length (100 um tolerance) - if isinstance( - ee_test, (pysixtrack.elements.Drift, pysixtrack.elements.DriftExact) - ): - if kk == "length": - if diff_abs < 1e-4: - continue - - # Exception: multipole lrad is not passed to sixtraxk - if isinstance(ee_test, pysixtrack.elements.Multipole): - if kk == "length": - if np.abs(ee_test.hxl) + np.abs(ee_test.hyl) == 0.0: - continue - - # Exceptions BB4D (separations are recalculated) - if isinstance(ee_test, pysixtrack.elements.BeamBeam4D): - if kk == "x_bb": - if diff_abs / dtest["sigma_x"] < 0.01: # This is neede to accommodate different leveling routines (1% difference) - continue - if kk == "y_bb": - if diff_abs / dtest["sigma_y"] < 0.01: - continue - if kk == "sigma_x": - if diff_rel < 1e-5: - continue - if kk == "sigma_y": - if diff_rel < 1e-5: - continue - - # Exceptions BB4D (angles and separations are recalculated) - if isinstance(ee_test, pysixtrack.elements.BeamBeam6D): - if kk == "alpha": - if diff_abs < 10e-6: - continue - if kk == "x_bb_co": - if diff_abs / np.sqrt(dtest["sigma_11"]) < 0.015: - continue - if kk == "y_bb_co": - if diff_abs / np.sqrt(dtest["sigma_33"]) < 0.015: - continue - - # If it got here it means that no condition above is met - raise ValueError("Too large discrepancy!") -print( - """ -******************************************************************* - The line from test seq. and the line from reference are identical! -******************************************************************* -""" -) 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 new file mode 100644 index 0000000..d6f366e --- /dev/null +++ b/python_examples/hl_lhc_collisions_python/checks_and_doc/t004_compare_xlines.py @@ -0,0 +1,276 @@ +import time +import shutil +import pickle +import numpy as np + +import xline +import sixtracktools + + +# Tests b1 with bb +tests = [ + { + 'test_name': 'B1 - pymask sixtrack input vs madx mask', + 'path_test': '../', + 'type_test': 'sixtrack', + 'path_ref': '../../../examples/hl_lhc_collision', + 'type_ref': 'sixtrack', + 'rtol': 3e-5, + 'atol': 1e-12, + 'strict': False, + }, + { + 'test_name': 'B1 - pymask xline vs pymask sixtrack input', + 'path_test': '../xlines/line_bb_dipole_not_cancelled.json', + 'type_test': 'xline', + 'path_ref': '../', + 'type_ref': 'sixtrack', + 'rtol': 4e-7, + 'atol': 1e-100, + 'strict': True, + } +] + +# # Tests b4 no bb +# tests = [ +# { +# 'test_name': 'B4 - pymask sixtrack input vs madx mask', +# 'path_test': '../', +# 'type_test': 'sixtrack', +# 'path_ref': '../../../examples/hl_lhc_collision_nobb_b4', +# 'type_ref': 'sixtrack', +# 'rtol': 8e-5, +# 'atol': 1e-12, +# 'strict': False, +# }, +# { +# 'test_name': 'B4 - pymask xline vs pymask sixtrack input', +# 'path_test': '../xline/line_bb_dipole_not_cancelled.json', +# 'type_test': 'xline', +# 'path_ref': '../', +# 'type_ref': 'sixtrack', +# 'rtol': 4e-7, +# 'atol': 1e-100, +# 'strict': True, +# } +# ] + +def norm(x): + return np.sqrt(np.sum(np.array(x) ** 2)) + +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) + elif input_type == 'sixtrack': + print('Build xline from sixtrack input:') + sixinput_test = sixtracktools.sixinput.SixInput(path) + ltest = xline.Line.from_sixinput(sixinput_test) + print('Done') + else: + raise ValueError('What?!') + + return ltest + +for tt in tests: + test_name = tt['test_name'] + path_test = tt['path_test'] + type_test = tt['type_test'] + path_ref = tt['path_ref'] + type_ref = tt['type_ref'] + rtol = tt['rtol'] + atol = tt['atol'] + strict = tt['strict'] + + # Load + ltest = prepare_line(path_test, type_test) + lref = prepare_line(path_ref, type_ref) + + original_length = ltest.get_length() + assert (lref.get_length() - original_length) < 1e-6 + + # Simplify the two machines + for ll in (ltest, lref): + ll.remove_inactive_multipoles(inplace=True) + ll.remove_zero_length_drifts(inplace=True) + ll.merge_consecutive_drifts(inplace=True) + ll.merge_consecutive_multipoles(inplace=True) + + # Remove inactive RFMultipoles and normalize phase + for ii, ee in enumerate(ll.elements): + if ee.__class__.__name__ == 'RFMultipole': + if ( np.max(np.abs(ee.knl)) < 1e-20 and + np.max(np.abs(ee.ksl)) < 1e-20 and + abs(ee.voltage) < 1e-20): + ll.element_names[ii] = None + ll.elements[ii] = None + # # Normalize phase + # for kkll, pp in [[ee.knl, ee.pn], + # [ee.ksl, ee.ps]]: + # for ii, vv in enumerate(kkll): + # if vv < 0: + # kkll[ii] = -vv + # pp[ii] += 180 + # if pp[ii]>180: + # pp[ii] -= 360 + + while None in ll.element_names: + ll.element_names.remove(None) + while None in ll.elements: + ll.elements.remove(None) + + # Check that the two machines are identical + assert len(ltest) == len(lref) + + assert (ltest.get_length() - original_length) < 1e-6 + assert (lref.get_length() - original_length) < 1e-6 + + for ii, (ee_test, ee_six, nn_test, nn_six) in enumerate( + zip(ltest.elements, lref.elements, ltest.element_names, lref.element_names) + ): + assert type(ee_test) == type(ee_six) + + dtest = ee_test.to_dict(keepextra=True) + dref = ee_six.to_dict(keepextra=True) + + for kk in dtest.keys(): + + # Check if they are identical + if np.isscalar(dref[kk]) and dtest[kk] == dref[kk]: + continue + + # Check if the relative error is small + val_test = dtest[kk] + val_ref = dref[kk] + try: + if not np.isscalar(val_ref) and len(val_ref) != len(val_test): + diff_rel = 100 + #lmin = min(len(val_ref), len(val_test)) + #val_test = val_test[:lmin] + #val_ref = val_ref[:lmin] + else: + diff_rel = norm(np.array(val_test) - np.array(val_ref)) / norm(val_test) + except ZeroDivisionError: + diff_rel = 100.0 + if diff_rel < rtol: + continue + + # Check if absolute error is small + + if not np.isscalar(val_ref) and len(val_ref) != len(val_test): + diff_abs = 1000 + else: + diff_abs = norm(np.array(val_test) - np.array(val_ref)) + if diff_abs > 0: + print(f"{nn_test}[{kk}] - test:{dtest[kk]} six:{dref[kk]}") + if diff_abs < atol: + continue + + # Exception: drift length (100 um tolerance) + if not(strict) and isinstance( + ee_test, (xline.elements.Drift, xline.elements.DriftExact) + ): + 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 kk == "length": + if np.abs(ee_test.hxl) + np.abs(ee_test.hyl) == 0.0: + continue + if kk == 'knl' or kk == 'ksl': + 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') + + val_ref = val_ref[:lmin] + val_test = val_test[:lmin] + + if len(val_ref) == 0 and len(val_test) == 0: + continue + else: + diff_abs = norm(np.array(val_test) - np.array(val_ref)) + diff_rel = diff_abs/norm(val_test) + if diff_rel < rtol: + continue + if diff_abs < atol: + continue + + # Exception: correctors involved in lumi leveling + passed_corr = False + for nn_corr in [ + 'mcbcv.5l8.b1', 'mcbyv.a4l8.b1', 'mcbxv.3l8', + 'mcbyv.4r8.b1', 'mcbyv.b5r8.b1', + 'mcbyh.b5l2.b1', 'mcbyh.4l2.b1', 'mcbxh.3l2', 'mcbyh.a4r2.b1', + 'mcbch.5r2.b1', + 'mcbcv.5l8.b2', 'mcbyv.a4l8.b2', 'mcbxv.3l8', + 'mcbyv.4r8.b2', 'mcbyv.b5r8.b2', + 'mcbyh.b5l2.b2', 'mcbyh.4l2.b2', 'mcbxh.3l2', 'mcbyh.a4r2.b2', + 'mcbch.5r2.b2', 'mcbch.a5r2.b2', 'mcbyh.4r2.b2', 'mcbxh.3r2', + 'mcbyh.a4l2.b2', 'mcbyh.5l2.b2', 'mcbyv.5r8.b2', 'mcbyv.a4r8.b2', + 'mcbxv.3r8', 'mcbyv.4l8.b2', 'mcbcv.b5l8.b2']: + if nn_corr in nn_test and diff_rel < 1e-2: + passed_corr = True + break + if not(strict) and passed_corr: + continue + + # Exceptions BB4D (separations are recalculated) + if not(strict) and isinstance(ee_test, xline.elements.BeamBeam4D): + if kk == "x_bb": + if diff_abs / dtest["sigma_x"] < 0.01: # This is neede to accommodate different leveling routines (1% difference) + continue + if kk == "y_bb": + if diff_abs / dtest["sigma_y"] < 0.01: + continue + if kk == "sigma_x": + if diff_rel < 1e-5: + continue + if kk == "sigma_y": + if diff_rel < 1e-5: + continue + + # Exceptions BB6D (angles and separations are recalculated) + if not(strict) and isinstance(ee_test, xline.elements.BeamBeam6D): + if kk == "alpha": + if diff_abs < 10e-6: + continue + if kk == "x_co" or kk == "x_bb_co": + if diff_abs / np.sqrt(dtest["sigma_11"]) < 0.015: + continue + if kk == "y_co" or kk == "y_bb_co": + if diff_abs / np.sqrt(dtest["sigma_33"]) < 0.015: + continue + if kk == "zeta_co": + if diff_abs <1e-5: + continue + if kk == "delta_co": + if diff_abs <1e-5: + continue + if kk == "px_co" or kk == 'py_co': + if diff_abs <30e-9: + continue + + # If it got here it means that no condition above is met + raise ValueError("Too large discrepancy!") + print( + f""" + ******************************************************************* + test: {test_name} + + The line from test seq. and the line from reference are identical! + ******************************************************************* + """ + ) + + time.sleep(1) 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 b4eab61..8596a70 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 @@ -1,8 +1,8 @@ import shutil -import pickle import numpy as np +import json -import pysixtrack +import xline import sixtracktools L_lhc = 27e3 @@ -46,16 +46,18 @@ path_test = '../' type_test = 'sixtrack' +path_test = '../xlines/line_bb_dipole_not_cancelled.json' +type_test = 'xline' + def prepare_line(path, input_type): - if input_type == 'pysixtrack': - # Load pysixtrack machine - with open(path, "rb") as fid: - ltest = pysixtrack.Line.from_dict(pickle.load(fid)) + if input_type == 'xline': + # Load xline machine + ltest = xline.Line.from_json(path) elif input_type == 'sixtrack': - print('Build pysixtrack from sixtrack input:') + print('Build xline from sixtrack input:') sixinput_test = sixtracktools.sixinput.SixInput(path) - ltest = pysixtrack.Line.from_sixinput(sixinput_test) + ltest = xline.Line.from_sixinput(sixinput_test) print('Done') else: raise ValueError('What?!') @@ -115,20 +117,14 @@ def prepare_line(path, input_type): # Chek crabs in weak beam -import pickle -with open('../optics_orbit_at_start_ring.pkl', 'rb') as fid: - ddd = pickle.load(fid) - -ddd['p0c'] = ddd['p0c_eV'] - # Switch off all beam-beam lenses -bb_all, _ = ltest.get_elements_of_type([pysixtrack.elements.BeamBeam4D, - pysixtrack.elements.BeamBeam6D]) +bb_all, _ = ltest.get_elements_of_type([xline.elements.BeamBeam4D, + xline.elements.BeamBeam6D]) for bb in bb_all: bb.enabled = False # # Switch off all beam-beam lenses -crabs, crab_names = ltest.get_elements_of_type([pysixtrack.elements.RFMultipole]) +crabs, crab_names = ltest.get_elements_of_type([xline.elements.RFMultipole]) #for cc in crabs: # cc.pn = [-90] # cc.ps = [-90] @@ -136,7 +132,9 @@ def prepare_line(path, input_type): # for cc in crabs: # cc.knl[0] *= -1 -partco = pysixtrack.Particles.from_dict(ddd) +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']) z_slices = s_rel * 2.0 partco.zeta += z_slices partco.x += 0*z_slices @@ -184,7 +182,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 pysixtrack') +axcrab.plot(s_rel, r_lenses, 'o', color='b', alpha=.5, label= 'weak xline') 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', @@ -222,7 +220,7 @@ def prepare_line(path, input_type): axcbx.plot(s_twiss, x_twiss) axcby.plot(s_twiss, y_twiss) -part = pysixtrack.Particles.from_dict(ddd) +part = xline.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]]) 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 6f184b6..a77029d 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 @@ -1,19 +1,15 @@ -import pickle +import json import numpy as np import NAFFlib import helpers as hp import footprint import matplotlib.pyplot as plt -import pysixtrack +import xline import sixtracktools -track_with = 'PySixtrack' -track_with = 'Sixtracklib' -#device = 'opencl:1.0' -device = None - -opt_start_ring_fname = '../optics_orbit_at_start_ring.pkl' +track_with = 'xtrack' +#track_with = 'sixtrack' epsn_x = 2.5e-6 epsn_y = 2.5e-6 @@ -25,40 +21,31 @@ def prepare_line(path, input_type): - if input_type == 'pysixtrack': - # Load pysixtrack machine - with open(path, "rb") as fid: - ltest = pysixtrack.Line.from_dict( - pickle.load(fid)) + if input_type == 'xline': + # Load xline machine + ltest = xline.Line.from_json(path) elif input_type == 'sixtrack': - print('Build pysixtrack from sixtrack input:') + print('Build xline from sixtrack input:') sixinput_test = sixtracktools.sixinput.SixInput(path) - ltest = pysixtrack.Line.from_sixinput(sixinput_test) + ltest = xline.Line.from_sixinput(sixinput_test) print('Done') else: raise ValueError('What?!') return ltest -line = prepare_line('../', input_type='sixtrack') +line = prepare_line('../xlines/line_bb_for_tracking.json', input_type='xline') -with open('../optics_orbit_at_start_ring.pkl', 'rb') as fid: - ddd = pickle.load(fid) -ddd['p0c'] = ddd['p0c_eV'] +with open('../xlines/line_bb_for_tracking.json', 'r') as fid: + partCO = xline.Particles.from_dict( + json.load(fid)['particle_on_tracker_co']) -partCO = pysixtrack.Particles.from_dict(ddd) +with open('../optics_orbit_at_start_ring_from_madx.json', 'r') as fid: + ddd = json.load(fid) -# line.disable_beambeam() part = partCO.copy() -line.beambeam_store_closed_orbit_and_dipolar_kicks( - partCO.copy(), - separation_given_wrt_closed_orbit_4D=True, - separation_given_wrt_closed_orbit_6D=True) - - - beta_x = ddd['betx'] beta_y = ddd['bety'] @@ -78,35 +65,30 @@ 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 == 'PySixtrack': +if track_with == 'xline': part = partCO.copy() - x_tbt, px_tbt, y_tbt, py_tbt, sigma_tbt, delta_tbt = hp.track_particle_pysixtrack( + 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 = hp.track_particle_sixtrack( +elif track_with == 'sixtrack': + x_tbt, px_tbt, y_tbt, py_tbt, sigma_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, - input_folder=sixtrack_input_folder) + input_folder='../') info = track_with - -elif track_with == 'Sixtracklib': - x_tbt, px_tbt, y_tbt, py_tbt, sigma_tbt, delta_tbt = hp.track_particle_sixtracklib( +elif track_with == 'xtrack': + x_tbt, px_tbt, y_tbt, py_tbt, sigma_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, device=device) + Dsigma_wrt_CO_m=0., Ddelta_wrt_CO=0., n_turns=n_turns) info = track_with - if device is None: - info += ' (CPU)' - else: - info += ' (GPU %s)'%device 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 new file mode 100644 index 0000000..bd74b32 --- /dev/null +++ b/python_examples/hl_lhc_collisions_python/checks_and_doc/t007_check_orbit_and_lin_normal_form.py @@ -0,0 +1,66 @@ +import json + +import numpy as np + +import xline as xl +import xtrack as xt + +with open('../xlines/line_bb_for_tracking.json', 'r') as fid: + line_dict = json.load(fid) + +line = xl.Line.from_dict(line_dict) + +partCO = xl.Particles.from_dict(line_dict['particle_on_tracker_co']) + +tracker = xt.Tracker(sequence=line) +particles = xt.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) + +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. +norm_emit_x = 2.5e-6 +geom_emit_x = norm_emit_x / particles.beta0 / particles.gamma0 + +n_part = 100 +theta = np.linspace(0, 2*np.pi, n_part) +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()) + + +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.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 new file mode 100644 index 0000000..490c01c --- /dev/null +++ b/python_examples/hl_lhc_collisions_python/checks_and_doc/t008_check_against_sixtrack.py @@ -0,0 +1,58 @@ +import json +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 sixtracktools + + +displace_x = [-5e-4, 5e-4] +displace_y = [3e-4, -3e-4] +num_turns = 5 + + +with open('../xlines/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) + +partCO = xl.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( + 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, + 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) + +particles = xt.Particles(**part_track.to_dict()) +particles.particle_id = np.arange(particles.num_particles) + +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) + +assert np.allclose(tracker.record_last_track.x[0, :], x_tbt_sixtrack[:,0], + rtol=1e-15, atol=9e-11) +assert np.allclose(tracker.record_last_track.y[0, :], y_tbt_sixtrack[:,0], + rtol=1e-15, atol=9e-11) +assert np.allclose(tracker.record_last_track.delta[0, :], delta_tbt_sixtrack[:,0], + rtol=1e-15, atol=5e-11) 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 ee593a1..6e97704 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 @@ -1,19 +1,14 @@ # Tests performed on pymask -Sixtracktools is on the master - -We test the commit 5da242cf332661f06a2151dc49fbf616b9806d6b - N.B. For the tests to be successful make sure that cpymad and madx correspond to the same (updated!) version. - ## Test 1 - b1 with bb legacy macros We execute the reference: ```bash cd ../../examples/hl_lhc_collision - python ../../unmask.py main.mask parameters_for_unmask.txt --run + python ../../unmask.py main.mask parameters_for_unmask.txt --run-cpymad ``` We setup the python version: @@ -54,22 +49,40 @@ python 000_pymask.py We execute the reference (can be reused from the previous test): ```bash cd ../../examples/hl_lhc_collision - python ../../unmask.py main.mask parameters_for_unmask.txt --run + python ../../unmask.py main.mask parameters_for_unmask.txt --run-cpymad ``` In the folder ```checks_and_doc```: -Configure ```t004_compare_pysixtrack_lines.py``` for the test: +Configure ```t004_compare_xline_lines.py``` for the test: ```python # Test b1 -path_test = '../' -type_test = 'sixtrack' -path_ref = '../../examples/hl_lhc_collision' -type_ref = 'sixtrack -``` - -Check using pysixtrack lines: +tests = [ + { + 'test_name': 'B1 - pymask sixtrack input vs madx mask', + 'path_test': '../', + 'type_test': 'sixtrack', + 'path_ref': '../../../examples/hl_lhc_collision', + 'type_ref': 'sixtrack', + 'rtol': 3e-5, + 'atol': 1e-12, + 'strict': False, + }, + { + 'test_name': 'B1 - pymask xline vs pymask sixtrack input', + 'path_test': '../xline/line_bb_dipole_not_cancelled.json', + 'type_test': 'xline', + 'path_ref': '../', + 'type_ref': 'sixtrack', + 'rtol': 4e-7, + 'atol': 1e-100, + 'strict': True, + } +] +``` + +Check using xline lines: ```bash python t003_fc_to_fort.py -python t004_compare_pysixtrack_lines.py +python t004_compare_xlines.py.py ``` **Check crab cavities:** @@ -100,7 +113,7 @@ Repeat for IP1. We execute the reference: ```bash cd ../../examples/hl_lhc_collision_nobb_b4 - python ../../unmask.py main.mask parameters_for_unmask.txt --run + python ../../unmask.py main.mask parameters_for_unmask.txt --run-cpymad ``` We setup the python version: @@ -134,7 +147,7 @@ diff fc.34 ../../examples/hl_lhc_collision_nobb_b4/fc.34 We execute the reference: ```bash cd ../../examples/hl_lhc_collision_nobb_b4 - python ../../unmask.py main.mask parameters_for_unmask.txt --run + python ../../unmask.py main.mask parameters_for_unmask.txt --run-cpymad ``` We setup the python version: diff --git a/python_examples/hl_lhc_collisions_python/config.py b/python_examples/hl_lhc_collisions_python/config.py index 3a85ea3..ba45642 100644 --- a/python_examples/hl_lhc_collisions_python/config.py +++ b/python_examples/hl_lhc_collisions_python/config.py @@ -7,6 +7,7 @@ 'tools': 'tracking_tools/tools', 'beambeam_macros': 'tracking_tools/beambeam_macros', 'errors': 'tracking_tools/errors', + 'optics_repository' : '/afs/cern.ch/eng/lhc/optics', }, # Mode - choose between: @@ -20,13 +21,13 @@ # 'b1_with_bb_legacy_macros' # 'b4_without_bb' - 'mode' : 'b1_with_bb_legacy_macros', + 'mode' : 'b1_with_bb', # For test against madx mask for modes 'b4_from_b2_without_bb' and 'b4_without_bb': - # 'force_leveling' : {'on_sep8': -0.03425547139366354, 'on_sep2': 0.14471680504084292}, + #'force_leveling' : {'on_sep8': -0.03425547139366354, 'on_sep2': 0.14471680504084292}, # Optics file - 'optics_file' : '/afs/cern.ch/eng/lhc/optics/HLLHCV1.4/round/opt_round_150_1500_thin.madx', #15 cm + 'optics_file' : 'optics_repository/HLLHCV1.4/round/opt_round_150_1500_thin.madx', #15 cm # Enable checks 'check_betas_at_ips' : True, @@ -72,13 +73,16 @@ 'nco_IP8' : 2572, # number of Head-On collisions in IP1 # Beam-beam parameters (used by python tools - NOT by legacy macros) - 'numberOfLRPerIRSide' : [25, 20, 25, 20], - 'bunch_spacing_buckets' : 10, - 'numberOfHOSlices' : 11, - 'bunch_population_ppb' : None, - 'sigmaz_m' : None, - 'z_crab_twiss' : 0.075, - + 'beambeam_config' : + { + 'numberOfLRPerIRSide' : [25, 20, 25, 20], + 'bunch_spacing_buckets' : 10, + 'numberOfHOSlices' : 11, + 'bunch_num_particles' : None, + 'bunch_particle_charge' : None, + 'sigmaz_m' : None, + 'z_crab_twiss' : 0.075, + }, # Match tunes and chromaticities including beam-beam effects 'match_q_dq_with_bb' : False, # should be off at collision diff --git a/python_examples/hl_lhc_collisions_python/customization.bash.example b/python_examples/hl_lhc_collisions_python/customization.bash.example new file mode 100644 index 0000000..3ef2c56 --- /dev/null +++ b/python_examples/hl_lhc_collisions_python/customization.bash.example @@ -0,0 +1,5 @@ +#rm modules +#ln -fns /home/sterbini/test_new_mask/v0.5/lhcmachines machines +#ln -fns /home/sterbini/test_new_mask/v0.5/lhctoolkit tools +ln -fns /home/sterbini/lhcmask modules +#ln -fns /home/sterbini/test_new_mask/v0.5/lhcerrors errors diff --git a/python_examples/hl_lhc_collisions_python/optics_specific_tools.py b/python_examples/hl_lhc_collisions_python/optics_specific_tools.py index f4f3cea..ca99362 100644 --- a/python_examples/hl_lhc_collisions_python/optics_specific_tools.py +++ b/python_examples/hl_lhc_collisions_python/optics_specific_tools.py @@ -11,7 +11,7 @@ def build_sequence(mad, beam): mad.input(f'mylhcbeam = {beam}') # Make link to optics toolkit - pm.make_links({'optics_toolkit': '/afs/cern.ch/eng/lhc/optics/HLLHCV1.4/toolkit'}, + pm.make_links({'optics_toolkit': 'optics_repository/HLLHCV1.4/toolkit'}, force=True) mad.input(''' @@ -21,20 +21,20 @@ def build_sequence(mad, beam): ver_hllhc_optics = 1.4; ! Get the toolkit - call, file="/afs/cern.ch/eng/lhc/optics/HLLHCV1.4/toolkit/macro.madx"; + call, file="optics_repository/HLLHCV1.4/toolkit/macro.madx"; ! Build sequence option, -echo,-warn,-info; if (mylhcbeam==4){ - call,file="/afs/cern.ch/eng/lhc/optics/runIII/lhcb4.seq"; + call,file="optics_repository/runIII/lhcb4.seq"; } else { - call,file="/afs/cern.ch/eng/lhc/optics/runIII/lhc.seq"; + call,file="optics_repository/runIII/lhc.seq"; }; option, -echo, warn,-info; !Install HL-LHC - call, file="/afs/cern.ch/eng/lhc/optics/HLLHCV1.4/hllhc_sequence.madx"; + call, file="optics_repository/HLLHCV1.4/hllhc_sequence.madx"; ! Slice nominal sequence 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 new file mode 100644 index 0000000..21e6f32 --- /dev/null +++ b/python_examples/hl_lhc_collisions_python/t002_test_save_lines.py @@ -0,0 +1,40 @@ +# Requires dataframes to be saved manually after executing 000_pymask.py + +import sys +import json + +import numpy as np + +import xline as xl +import xtrack as xt +from scipy.optimize import fsolve + +sys.path.append('./modules') +import pymask as pm + +Madx = pm.Madxp +mad = Madx(command_log="mad_final.log") +mad.call("final_seq.madx") +mad.use(sequence="lhcb1") +mad.twiss() +mad.readtable(file="final_errors.tfs", table="errtab") +mad.seterr(table="errtab") +mad.set(format=".15g") +mad.twiss(rmatrix = True) + +separation_given_wrt_closed_orbit_4D = True +tw = mad.table['twiss'] + +seq_name = 'lhcb1' +mad_beam = mad.sequence[seq_name].beam + +import pandas as pd +bb_df_track = pd.read_pickle('./bb_df_track.pkl') +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, + optics_and_co_at_start_ring_from_madx, + folder_name = 'xlines', + 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 0f37d14..76dc68c 100644 --- a/python_examples/ions_python/000_pymask.py +++ b/python_examples/ions_python/000_pymask.py @@ -1,6 +1,6 @@ import os import sys -import pickle +import json import numpy as np @@ -89,19 +89,19 @@ mad.globals.nrj = configuration['beam_energy_tot'] particle_type = 'proton' -if 'beam_mass' in configuration.keys(): - beam_mass = configuration['beam_mass'] +if 'particle_mass' in configuration.keys(): + particle_mass = configuration['particle_mass'] particle_type = 'ion' else: - beam_mass = mad.globals.pmass # proton mass + particle_mass = mad.globals.pmass # proton mass -if 'beam_charge' in configuration.keys(): - beam_charge = configuration['beam_charge'] +if 'particle_charge' in configuration.keys(): + particle_charge = configuration['particle_charge'] particle_type = 'ion' else: - beam_charge = 1. + particle_charge = 1. -gamma_rel = (beam_charge*configuration['beam_energy_tot'])/beam_mass +gamma_rel = (particle_charge*configuration['beam_energy_tot'])/particle_mass for ss in mad.sequence.keys(): # bv and bv_aux flags if ss == 'lhcb1': @@ -115,15 +115,15 @@ mad.globals['bv_aux'] = ss_bv_aux mad.input(f''' beam, particle={particle_type},sequence={ss}, - energy={configuration['beam_energy_tot']*beam_charge}, + energy={configuration['beam_energy_tot']*particle_charge}, sigt={configuration['beam_sigt']}, bv={ss_beam_bv}, npart={configuration['beam_npart']}, sige={configuration['beam_sige']}, ex={configuration['beam_norm_emit_x'] * 1e-6 / gamma_rel}, ey={configuration['beam_norm_emit_y'] * 1e-6 / gamma_rel}, - mass={beam_mass}, - charge={beam_charge}, + mass={particle_mass}, + charge={particle_charge}, ''') @@ -247,20 +247,21 @@ # Prepare bb dataframes if enable_bb_python: + bbconfig = configuration['beambeam_config'] bb_dfs = pm.generate_bb_dataframes(mad, ip_names=['ip1', 'ip2', 'ip5', 'ip8'], harmonic_number=35640, - numberOfLRPerIRSide=configuration['numberOfLRPerIRSide'], - bunch_spacing_buckets=configuration['bunch_spacing_buckets'], - numberOfHOSlices=configuration['numberOfHOSlices'], - bunch_population_ppb=configuration['bunch_population_ppb'], - sigmaz_m=configuration['sigmaz_m'], - z_crab_twiss=configuration['z_crab_twiss']*float(enable_crabs), + numberOfLRPerIRSide=bbconfig['numberOfLRPerIRSide'], + bunch_spacing_buckets=bbconfig['bunch_spacing_buckets'], + numberOfHOSlices=bbconfig['numberOfHOSlices'], + bunch_num_particles = bbconfig['bunch_num_particles'], + bunch_particle_charge = bbconfig['bunch_particle_charge'], + sigmaz_m=bbconfig['sigmaz_m'], + z_crab_twiss=bbconfig['z_crab_twiss']*float(enable_crabs), remove_dummy_lenses=True) # Here the datafremes can be edited, e.g. to set bbb intensity - ################### # Generate beam 4 # ################### @@ -286,11 +287,12 @@ check_betas_at_ips=check_betas_at_ips, check_separations_at_ips=False) # For B1, to be generalized for B4 -if 'filling_scheme_json' in configuration.keys(): - filling_scheme_json = configuration['filling_scheme_json'] - bunch_to_track = configuration['bunch_to_track'] +if 'filling_scheme_json' in configuration['beambeam_config'].keys(): + assert 'b4' not in mode + filling_scheme_json = configuration['beambeam_config']['filling_scheme_json'] + bunch_to_track = configuration['beambeam_config']['bunch_to_track'] bb_schedule_to_track_b1 = ost.create_bb_shedule_to_track( - filling_scheme_json,bunch_to_track, beam='1') + filling_scheme_json,bunch_to_track, beam=1) bb_dfs['b1']=ost.filter_bb_df(bb_dfs['b1'],bb_schedule_to_track_b1) ################################################## @@ -337,6 +339,7 @@ # Legacy bb macros if enable_bb_legacy: + bbconfig = configuration['beambeam_config'] assert(beam_to_configure == 1) assert(not(track_from_b4_mad_instance)) assert(not(enable_bb_python)) @@ -344,7 +347,7 @@ mad_track.set_variables_from_dict( params=configuration['pars_for_legacy_bb_macros']) mad_track.set_variables_from_dict( - params={f'par_nho_ir{ir}':configuration['numberOfHOSlices'] + params={f'par_nho_ir{ir}': bbconfig['numberOfHOSlices'] for ir in [1,2,5,8]}) mad_track.input("call, file='modules/module_03_beambeam.madx';") @@ -492,8 +495,9 @@ seq_name=sequence_to_track, bb_df=bb_df_track, output_folder='./', - reference_bunch_charge_sixtrack_ppb=( + reference_num_particles_sixtrack=( mad_track.sequence[sequence_to_track].beam.npart), + reference_particle_charge_sixtrack=mad_track.sequence[sequence_to_track].beam.charge, emitnx_sixtrack_um=( mad_track.sequence[sequence_to_track].beam.exn), emitny_sixtrack_um=( @@ -514,29 +518,27 @@ # Save optics and orbit at start ring # ####################################### -optics_orbit_start_ring = pm.get_optics_and_orbit_at_start_ring( +optics_and_co_at_start_ring_from_madx = pm.get_optics_and_orbit_at_start_ring( mad_track, sequence_to_track, skip_mad_use=True) -with open('./optics_orbit_at_start_ring.pkl', 'wb') as fid: - pickle.dump(optics_orbit_start_ring, fid) - +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 pysixtrack lines # -############################# +################## +# Generate xline # +################## if enable_bb_legacy: - print('Pysixtrack line is not generated with bb legacy macros') + print('xline is not generated with bb legacy macros') else: - pysix_fol_name = "./pysixtrack" - dct_pysxt = pm.generate_pysixtrack_line_with_bb(mad_track, - sequence_to_track, bb_df_track, - closed_orbit_method='from_mad', - pickle_lines_in_folder=pysix_fol_name, - skip_mad_use=True) + pm.generate_xline(mad_track, sequence_to_track, bb_df_track, + optics_and_co_at_start_ring_from_madx, + folder_name = './xlines', + skip_mad_use=True, + prepare_line_for_xtrack=True) -############################# -# Save final twiss # -############################# +################################### +# Save final twiss # +################################### mad_track.globals.on_bb_charge = 0 mad_track.twiss() @@ -544,3 +546,17 @@ sdf = mad_track.get_summ_df('summ') tdf.to_parquet('final_twiss_BBOFF.parquet') sdf.to_parquet('final_summ_BBOFF.parquet') + + +mad_track.globals.on_bb_charge = 1 +mad_track.twiss() +tdf = mad_track.get_twiss_df('twiss') +sdf = mad_track.get_summ_df('summ') +tdf.to_parquet('final_twiss_BBON.parquet') +sdf.to_parquet('final_summ_BBON.parquet') + +############################# +# Save sequence and errors # +############################# +# N.B. this erases the errors in the mad_track instance +# pm.save_mad_sequence_and_error(mad_track, sequence_to_track, filename='final') diff --git a/python_examples/ions_python/checks_and_doc/fort_parts/fort_beginning.3 b/python_examples/ions_python/checks_and_doc/fort_parts/fort_beginning.3 new file mode 100644 index 0000000..e48958c --- /dev/null +++ b/python_examples/ions_python/checks_and_doc/fort_parts/fort_beginning.3 @@ -0,0 +1,36 @@ +GEOME-STRENG TITLE:ts_collisions_ats30_en20_IMO380_C7_X160_I1.2_62.3100_60.3200%1%s%62.31_60.32%8_10%5%80 +PRINTOUT OF INPUT PARAMETERS-------------------------------------------- +NEXT +TRACKING PARAMETERS----------------------------------------------------- +1 0 1 0.00 0.00 0 1 +1 1 0 1 2 +0 0 1 1 1000 50000 2 +NEXT +INITIAL COORDINATES----------------------------------------------------- + 2 0. 0. 0.999999 0 + 0. + 0. + 0. + 0. + 0. + 0. + .0 + 0.01 + 0.003 + 0. + 0. + 0.000 + 574000000.0 + 574000000.0 + 574000000.0 +NEXT +SYNCHROTRON OSCILLATIONS--------'PLACE AFTER TRACKING PARAMETERS'------- + 35640 .000347 14. 0. 26658.883200 193687.2729 1 + 1. 1. +NEXT +HION +208 82 193.6872729 82 +NEXT +DUMP +ALL 1 664 101 dump3.dat +NEXT diff --git a/python_examples/ions_python/checks_and_doc/fort_parts/fort_end.3 b/python_examples/ions_python/checks_and_doc/fort_parts/fort_end.3 new file mode 100644 index 0000000..64968ba --- /dev/null +++ b/python_examples/ions_python/checks_and_doc/fort_parts/fort_end.3 @@ -0,0 +1,57 @@ +FLUCTUATION +100000 1 7 3 +NEXT +/TUNE VARIATION---------------------------------------------------------- +/ mq.7r3.b1 62.31 +/ mq.8r3.b1 60.32 +// 0. +// +/NEXT +/CHROMATICITY CORRECTION------------------------------------------------ +/ ms.11r3.b1 2. +/ ms.12r3.b1 2. +/NEXT +POSTPROCESSING---------------------------------------------------------- +LHC Dynamic Aperture at Injection Version 5 (1998) +20 0 0 1 .08 .08 1 0 0 1. 1. +62. 60. 1 1 10 .005 1 .050 +0 1 0 1 1 1 1 1 30 +NEXT +/DIFFERENTIAL ALGEBRA--------------------------------------------------- +/%NO %NV 1E-38 %nsix 0 +/NEXT +ENDE==================================================================== +ORGA-------------------------------------------------------------------- +NEXT +DIFFERENTIAL ALGEBRA--------------------------------------------------- +4 5 1E-38 1 1 +NEXT +COMBINATION OF ELEMENTS------------------------------------------------- +CODSDO 1. CODSDE +CODSFO 1. CODSFE +NEXT +SUBRESONANCE CALCULATION------------------------------------------------ +3 3 62.31 60.32 1. 1. 0 26658.883200 +NEXT +RESONANCE COMPENSATION ------------------------------------------------- + 1 3 0 0 0 0 0 0 + 0 0 0 0 + 26658.883200 62.31 60.32 1. 1. + MSCBH MSCBV + 0 MSCBH MSCBV + 0 QF QD 62.31 60.32 +NEXT +RIPPLE OF POWER SUPPLIES------------------------------------------------ + Q413R -1.54D-4 4340.580 + Q414R -3.70D-4 4340.580 +NEXT +LIMITATION OF APERTURE-------------------------------------------------- + AP1 EL 10. 10. +NEXT +SEARCH FOR THE BEST POSITIONS FOR COMPENSATION-ELEMENTS----------------- + 62.31 60.32 1. 1. 26658.883200 + 2 3 0 2 -2 0 0 0 + MSCBH MSCBV +NEXT +ORBIT ADJUSTMENT-------------------------------------------------------- +NEXT diff --git a/python_examples/ions_python/checks_and_doc/helpers.py b/python_examples/ions_python/checks_and_doc/helpers.py new file mode 100644 index 0000000..8caf3b0 --- /dev/null +++ b/python_examples/ions_python/checks_and_doc/helpers.py @@ -0,0 +1,271 @@ +import numpy as np +import os +import sixtracktools +import xline +import time + + +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): + + 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]: + if hasattr(var, '__iter__'): + if n_part == 1: + n_part = len(var) + assert len(var) == n_part + + Dx_wrt_CO_m = Dx_wrt_CO_m + np.zeros(n_part) + 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) + 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 + + +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, + 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( + Dx_wrt_CO_m, Dpx_wrt_CO_rad, + Dy_wrt_CO_m, Dpy_wrt_CO_rad, + Dsigma_wrt_CO_m, Ddelta_wrt_CO) + + n_part = len(Dx_wrt_CO_m) + + wfold = 'temp_trackfun' + + if not os.path.exists(wfold): + os.mkdir(wfold) + + os.system(f'cp {input_folder}/fort.* {wfold}') + + with open(f'{wfold}/fort.3', 'r') as fid: + lines_f3 = fid.readlines() + + # Set initial coordinates + i_start_ini = None + for ii, ll in enumerate(lines_f3): + if ll.startswith('INITIAL COO'): + i_start_ini = ii + break + + lines_f3[i_start_ini + 2] = ' 0.\n' + lines_f3[i_start_ini + 3] = ' 0.\n' + lines_f3[i_start_ini + 4] = ' 0.\n' + lines_f3[i_start_ini + 5] = ' 0.\n' + lines_f3[i_start_ini + 6] = ' 0.\n' + lines_f3[i_start_ini + 7] = ' 0.\n' + + lines_f3[i_start_ini + 2 + 6] = ' 0.\n' + lines_f3[i_start_ini + 3 + 6] = ' 0.\n' + lines_f3[i_start_ini + 4 + 6] = ' 0.\n' + lines_f3[i_start_ini + 5 + 6] = ' 0.\n' + lines_f3[i_start_ini + 6 + 6] = ' 0.\n' + lines_f3[i_start_ini + 7 + 6] = ' 0.\n' + + lines_f13 = [] + + for i_part in range(n_part): + if type(partCO) is xline.Particles: + temp_part = partCO.copy() + else: + temp_part = xline.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] + + 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.delta))) + if i_part % 2 == 1: + lines_f13.append('%.10e\n' % (temp_part.energy0 * 1e-6)) + lines_f13.append('%.10e\n' % (prev_part.energy * 1e-6)) + lines_f13.append('%.10e\n' % (temp_part.energy * 1e-6)) + prev_part = temp_part + + with open(wfold + '/fort.13', 'w') as fid: + fid.writelines(lines_f13) + + if np.mod(n_part, 2) != 0: + raise ValueError('SixTrack does not like this!') + + i_start_tk = None + for ii, ll in enumerate(lines_f3): + if ll.startswith('TRACKING PAR'): + i_start_tk = ii + break + # Set number of turns and number of particles + temp_list = lines_f3[i_start_tk + 1].split(' ') + temp_list[0] = '%d' % n_turns + temp_list[2] = '%d' % (n_part / 2) + lines_f3[i_start_tk + 1] = ' '.join(temp_list) + # Set number of idfor = 2 + temp_list = lines_f3[i_start_tk + 2].split(' ') + temp_list[2] = '2' + lines_f3[i_start_tk + 2] = ' '.join(temp_list) + + # Setup turn-by-turn dump + i_start_dp = None + for ii, ll in enumerate(lines_f3): + if ll.startswith('DUMP'): + i_start_dp = ii + break + + lines_f3[i_start_dp + 1] = 'StartDUMP 1 664 101 dumtemp.dat\n' + + with open(wfold + '/fort.3', 'w') as fid: + fid.writelines(lines_f3) + + os.system('(cd temp_trackfun; sixtrack >fort.6)') + + # Load sixtrack tracking data + sixdump_all = sixtracktools.SixDump101('%s/dumtemp.dat' % wfold) + + x_tbt = np.zeros((n_turns, n_part)) + 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)) + delta_tbt = np.zeros((n_turns, n_part)) + + for i_part in range(n_part): + sixdump_part = sixdump_all[i_part::n_part] + x_tbt[:, i_part] = sixdump_part.x + 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 + delta_tbt[:, i_part] = sixdump_part.delta + + extra = {} + + return x_tbt, px_tbt, y_tbt, py_tbt, sigma_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, + _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( + Dx_wrt_CO_m, Dpx_wrt_CO_rad, + Dy_wrt_CO_m, Dpy_wrt_CO_rad, + Dsigma_wrt_CO_m, Ddelta_wrt_CO) + + import xtrack as xt + tracker = xt.Tracker(_context=_context, sequence=line) + particles = xt.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, + delta = partCO.delta + Ddelta_wrt_CO) + + print('Start track') + tracker.track(particles, num_turns=n_turns, turn_by_turn_monitor=True) + print('Done track') + + #print(res.particles[0]) + x_tbt = tracker.record_last_track.x.copy().T + 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 + 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 + + +def betafun_from_ellip(x_tbt, px_tbt): + + x_max = np.max(x_tbt) + mask = np.logical_and(np.abs(x_tbt) < x_max / 5., px_tbt > 0) + x_masked = x_tbt[mask] + px_masked = px_tbt[mask] + ind_sorted = np.argsort(x_masked) + x_sorted = np.take(x_masked, ind_sorted) + px_sorted = np.take(px_masked, ind_sorted) + + px_cut = np.interp(0, x_sorted, px_sorted) + + beta_x = x_max / px_cut + + return beta_x, x_max, px_cut diff --git a/python_examples/ions_python/checks_and_doc/t003_fc_to_fort.py b/python_examples/ions_python/checks_and_doc/t003_fc_to_fort.py new file mode 100644 index 0000000..6829953 --- /dev/null +++ b/python_examples/ions_python/checks_and_doc/t003_fc_to_fort.py @@ -0,0 +1,31 @@ +import os +import shutil + +folder_list = [ + '../', + '../../../examples/hl_lhc_collision/', + '../../../examples/hl_lhc_collision_nobb_b4'] + +for sixtrack_input_folder in folder_list: + print(sixtrack_input_folder) + try: + for iff in [2,8,16,34]: + os.system(f"rm {sixtrack_input_folder}/fort.{iff}") + try: + shutil.copy(sixtrack_input_folder + f"/fc.{iff}", + sixtrack_input_folder + f"/fort.{iff}") + except Exception: + print(f"/fc.{iff} not found!") + + with open(sixtrack_input_folder + "/fort.3", "w") as fout: + with open("fort_parts/fort_beginning.3", "r") as fid_fort3b: + fout.write(fid_fort3b.read()) + with open(sixtrack_input_folder + "/fc.3", "r") as fid_fc3: + fout.write(fid_fc3.read()) + with open("fort_parts/fort_end.3", "r") as fid_fort3e: + fout.write(fid_fort3e.read()) + except Exception as e: + print('Skipped...') + print(e) + +print('Ended.') 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 new file mode 100644 index 0000000..b01da6b --- /dev/null +++ b/python_examples/ions_python/checks_and_doc/t008_check_against_sixtrack.py @@ -0,0 +1,56 @@ +import json +import numpy as np +import helpers as hp +import matplotlib.pyplot as plt + +import xline as xl +import xtrack as xt +import sixtracktools + + +displace_x = [-5e-4, 5e-4] +displace_y = [3e-4, -3e-4] +num_turns = 5 + + +with open('../xlines/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) + +partCO = xl.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( + 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, + 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) + +particles = xt.Particles(**part_track.to_dict()) +particles.particle_id = np.arange(particles.num_particles) + +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) + +assert np.allclose(tracker.record_last_track.x[0, :], x_tbt_sixtrack[:,0], + rtol=1e-15, atol=9e-11) +assert np.allclose(tracker.record_last_track.y[0, :], y_tbt_sixtrack[:,0], + rtol=1e-15, atol=9e-11) +assert np.allclose(tracker.record_last_track.delta[0, :], delta_tbt_sixtrack[:,0], + rtol=1e-15, atol=5e-11) diff --git a/python_examples/ions_python/config.py b/python_examples/ions_python/config.py index daac18e..d97b213 100644 --- a/python_examples/ions_python/config.py +++ b/python_examples/ions_python/config.py @@ -3,11 +3,11 @@ # Links to be made for tools and scripts 'links' : { 'tracking_tools': '/afs/cern.ch/eng/tracking-tools', - #'modules': 'tracking_tools/modules', - 'modules': '../../', + 'modules': 'tracking_tools/modules', 'tools': 'tracking_tools/tools', 'beambeam_macros': 'tracking_tools/beambeam_macros', 'errors': 'tracking_tools/errors', + 'optics_repository' : '/afs/cern.ch/eng/lhc/optics', }, # Mode - choose between: @@ -24,7 +24,7 @@ 'mode' : 'b1_with_bb', # Optics file - 'optics_file' : '/afs/cern.ch/eng/lhc/optics/runII/2018/ION/opticsfile.21', + 'optics_file' : 'optics_repository/runII/2018/ION/opticsfile.21', # Enable checks 'check_betas_at_ips' : True, @@ -47,8 +47,8 @@ 'beam_energy_tot' : 7000. , # [GeV] # Ion parameters - 'beam_mass' :193.68715, - 'beam_charge' :82, + 'particle_mass' :193.6872729, + 'particle_charge' :82, # Tunes and chromaticities 'qx0' : 62.31, @@ -79,22 +79,27 @@ 'nco_IP8' : 398, # Beam-beam parameters (used by python tools - NOT by legacy macros) - 'numberOfLRPerIRSide' : [25, 20, 25, 20], - 'bunch_spacing_buckets' : 10, - 'numberOfHOSlices' : 11, - 'bunch_population_ppb' : None, - 'sigmaz_m' : None, - 'z_crab_twiss' : 0., - - 'bunch_to_track' : 488, - 'filling_scheme_json' : 'filling.json', + 'beambeam_config' : + { + 'numberOfLRPerIRSide' : [25, 20, 25, 20], + 'bunch_spacing_buckets' : 10, + 'numberOfHOSlices' : 11, + 'bunch_num_particles' : None, + 'bunch_particle_charge' : None, + 'sigmaz_m' : None, + 'z_crab_twiss' : 0., + + # Select from filling scheme + 'filling_scheme_json' : 'filling.json', + 'bunch_to_track' : 488, + }, # Match tunes and chromaticities including beam-beam effects 'match_q_dq_with_bb' : False, # should be off at collision # Enable crab cavities 'enable_crabs' : False, - # N. iterations coupling correction + # N. iterations for coupling correction 'N_iter_coupling' : 2, # Value to be added to linear coupling knobs (on sequence_to_track) diff --git a/python_examples/ions_python/customization.bash.example b/python_examples/ions_python/customization.bash.example new file mode 100644 index 0000000..3ef2c56 --- /dev/null +++ b/python_examples/ions_python/customization.bash.example @@ -0,0 +1,5 @@ +#rm modules +#ln -fns /home/sterbini/test_new_mask/v0.5/lhcmachines machines +#ln -fns /home/sterbini/test_new_mask/v0.5/lhctoolkit tools +ln -fns /home/sterbini/lhcmask modules +#ln -fns /home/sterbini/test_new_mask/v0.5/lhcerrors errors diff --git a/python_examples/run3_collisions_python/000_pymask.py b/python_examples/run3_collisions_python/000_pymask.py index 0f37d14..76dc68c 100644 --- a/python_examples/run3_collisions_python/000_pymask.py +++ b/python_examples/run3_collisions_python/000_pymask.py @@ -1,6 +1,6 @@ import os import sys -import pickle +import json import numpy as np @@ -89,19 +89,19 @@ mad.globals.nrj = configuration['beam_energy_tot'] particle_type = 'proton' -if 'beam_mass' in configuration.keys(): - beam_mass = configuration['beam_mass'] +if 'particle_mass' in configuration.keys(): + particle_mass = configuration['particle_mass'] particle_type = 'ion' else: - beam_mass = mad.globals.pmass # proton mass + particle_mass = mad.globals.pmass # proton mass -if 'beam_charge' in configuration.keys(): - beam_charge = configuration['beam_charge'] +if 'particle_charge' in configuration.keys(): + particle_charge = configuration['particle_charge'] particle_type = 'ion' else: - beam_charge = 1. + particle_charge = 1. -gamma_rel = (beam_charge*configuration['beam_energy_tot'])/beam_mass +gamma_rel = (particle_charge*configuration['beam_energy_tot'])/particle_mass for ss in mad.sequence.keys(): # bv and bv_aux flags if ss == 'lhcb1': @@ -115,15 +115,15 @@ mad.globals['bv_aux'] = ss_bv_aux mad.input(f''' beam, particle={particle_type},sequence={ss}, - energy={configuration['beam_energy_tot']*beam_charge}, + energy={configuration['beam_energy_tot']*particle_charge}, sigt={configuration['beam_sigt']}, bv={ss_beam_bv}, npart={configuration['beam_npart']}, sige={configuration['beam_sige']}, ex={configuration['beam_norm_emit_x'] * 1e-6 / gamma_rel}, ey={configuration['beam_norm_emit_y'] * 1e-6 / gamma_rel}, - mass={beam_mass}, - charge={beam_charge}, + mass={particle_mass}, + charge={particle_charge}, ''') @@ -247,20 +247,21 @@ # Prepare bb dataframes if enable_bb_python: + bbconfig = configuration['beambeam_config'] bb_dfs = pm.generate_bb_dataframes(mad, ip_names=['ip1', 'ip2', 'ip5', 'ip8'], harmonic_number=35640, - numberOfLRPerIRSide=configuration['numberOfLRPerIRSide'], - bunch_spacing_buckets=configuration['bunch_spacing_buckets'], - numberOfHOSlices=configuration['numberOfHOSlices'], - bunch_population_ppb=configuration['bunch_population_ppb'], - sigmaz_m=configuration['sigmaz_m'], - z_crab_twiss=configuration['z_crab_twiss']*float(enable_crabs), + numberOfLRPerIRSide=bbconfig['numberOfLRPerIRSide'], + bunch_spacing_buckets=bbconfig['bunch_spacing_buckets'], + numberOfHOSlices=bbconfig['numberOfHOSlices'], + bunch_num_particles = bbconfig['bunch_num_particles'], + bunch_particle_charge = bbconfig['bunch_particle_charge'], + sigmaz_m=bbconfig['sigmaz_m'], + z_crab_twiss=bbconfig['z_crab_twiss']*float(enable_crabs), remove_dummy_lenses=True) # Here the datafremes can be edited, e.g. to set bbb intensity - ################### # Generate beam 4 # ################### @@ -286,11 +287,12 @@ check_betas_at_ips=check_betas_at_ips, check_separations_at_ips=False) # For B1, to be generalized for B4 -if 'filling_scheme_json' in configuration.keys(): - filling_scheme_json = configuration['filling_scheme_json'] - bunch_to_track = configuration['bunch_to_track'] +if 'filling_scheme_json' in configuration['beambeam_config'].keys(): + assert 'b4' not in mode + filling_scheme_json = configuration['beambeam_config']['filling_scheme_json'] + bunch_to_track = configuration['beambeam_config']['bunch_to_track'] bb_schedule_to_track_b1 = ost.create_bb_shedule_to_track( - filling_scheme_json,bunch_to_track, beam='1') + filling_scheme_json,bunch_to_track, beam=1) bb_dfs['b1']=ost.filter_bb_df(bb_dfs['b1'],bb_schedule_to_track_b1) ################################################## @@ -337,6 +339,7 @@ # Legacy bb macros if enable_bb_legacy: + bbconfig = configuration['beambeam_config'] assert(beam_to_configure == 1) assert(not(track_from_b4_mad_instance)) assert(not(enable_bb_python)) @@ -344,7 +347,7 @@ mad_track.set_variables_from_dict( params=configuration['pars_for_legacy_bb_macros']) mad_track.set_variables_from_dict( - params={f'par_nho_ir{ir}':configuration['numberOfHOSlices'] + params={f'par_nho_ir{ir}': bbconfig['numberOfHOSlices'] for ir in [1,2,5,8]}) mad_track.input("call, file='modules/module_03_beambeam.madx';") @@ -492,8 +495,9 @@ seq_name=sequence_to_track, bb_df=bb_df_track, output_folder='./', - reference_bunch_charge_sixtrack_ppb=( + reference_num_particles_sixtrack=( mad_track.sequence[sequence_to_track].beam.npart), + reference_particle_charge_sixtrack=mad_track.sequence[sequence_to_track].beam.charge, emitnx_sixtrack_um=( mad_track.sequence[sequence_to_track].beam.exn), emitny_sixtrack_um=( @@ -514,29 +518,27 @@ # Save optics and orbit at start ring # ####################################### -optics_orbit_start_ring = pm.get_optics_and_orbit_at_start_ring( +optics_and_co_at_start_ring_from_madx = pm.get_optics_and_orbit_at_start_ring( mad_track, sequence_to_track, skip_mad_use=True) -with open('./optics_orbit_at_start_ring.pkl', 'wb') as fid: - pickle.dump(optics_orbit_start_ring, fid) - +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 pysixtrack lines # -############################# +################## +# Generate xline # +################## if enable_bb_legacy: - print('Pysixtrack line is not generated with bb legacy macros') + print('xline is not generated with bb legacy macros') else: - pysix_fol_name = "./pysixtrack" - dct_pysxt = pm.generate_pysixtrack_line_with_bb(mad_track, - sequence_to_track, bb_df_track, - closed_orbit_method='from_mad', - pickle_lines_in_folder=pysix_fol_name, - skip_mad_use=True) + pm.generate_xline(mad_track, sequence_to_track, bb_df_track, + optics_and_co_at_start_ring_from_madx, + folder_name = './xlines', + skip_mad_use=True, + prepare_line_for_xtrack=True) -############################# -# Save final twiss # -############################# +################################### +# Save final twiss # +################################### mad_track.globals.on_bb_charge = 0 mad_track.twiss() @@ -544,3 +546,17 @@ sdf = mad_track.get_summ_df('summ') tdf.to_parquet('final_twiss_BBOFF.parquet') sdf.to_parquet('final_summ_BBOFF.parquet') + + +mad_track.globals.on_bb_charge = 1 +mad_track.twiss() +tdf = mad_track.get_twiss_df('twiss') +sdf = mad_track.get_summ_df('summ') +tdf.to_parquet('final_twiss_BBON.parquet') +sdf.to_parquet('final_summ_BBON.parquet') + +############################# +# Save sequence and errors # +############################# +# N.B. this erases the errors in the mad_track instance +# pm.save_mad_sequence_and_error(mad_track, sequence_to_track, filename='final') diff --git a/python_examples/run3_collisions_python/config.py b/python_examples/run3_collisions_python/config.py index b954ae5..65c8233 100644 --- a/python_examples/run3_collisions_python/config.py +++ b/python_examples/run3_collisions_python/config.py @@ -7,6 +7,7 @@ 'tools': 'tracking_tools/tools', 'beambeam_macros': 'tracking_tools/beambeam_macros', 'errors': 'tracking_tools/errors', + 'optics_repository' : '/afs/cern.ch/eng/lhc/optics', }, # Mode - choose between: @@ -23,7 +24,7 @@ 'mode' : 'b1_with_bb', # Optics file - 'optics_file' : '/afs/cern.ch/eng/lhc/optics/runIII/RunIII_dev/2022_V1/PROTON/opticsfile.29', + 'optics_file' : 'optics_repository/runIII/RunIII_dev/2022_V1/PROTON/opticsfile.29', # Enable checks 'check_betas_at_ips' : True, @@ -69,12 +70,16 @@ 'nco_IP8' : 2376, # Beam-beam parameters (used by python tools - NOT by legacy macros) - 'numberOfLRPerIRSide' : [25, 20, 25, 20], - 'bunch_spacing_buckets' : 10, - 'numberOfHOSlices' : 11, - 'bunch_population_ppb' : None, - 'sigmaz_m' : None, - 'z_crab_twiss' : 0., + 'beambeam_config' : + { + 'numberOfLRPerIRSide' : [25, 20, 25, 20], + 'bunch_spacing_buckets' : 10, + 'numberOfHOSlices' : 11, + 'bunch_num_particles' : None, + 'bunch_particle_charge' : None, + 'sigmaz_m' : None, + 'z_crab_twiss' : 0., + }, # Match tunes and chromaticities including beam-beam effects 'match_q_dq_with_bb' : False, # should be off at collision @@ -82,7 +87,7 @@ # Enable crab cavities 'enable_crabs' : False, - # N. iterations coupling correction + # N. iterations for coupling correction 'N_iter_coupling' : 2, # Value to be added to linear coupling knobs (on sequence_to_track) diff --git a/python_examples/run3_collisions_python/old/000_pymask.py b/python_examples/run3_collisions_python/old/000_pymask.py deleted file mode 100644 index 411daab..0000000 --- a/python_examples/run3_collisions_python/old/000_pymask.py +++ /dev/null @@ -1,271 +0,0 @@ -import sys, os -import pickle - -import numpy as np - -# Import pymask -sys.path.append('../../') -import pymask as pm - -import optics_specific_tools as ost -from mask_parameters import mask_parameters - -Madx = pm.Madxp - -# Select mode -#mode = 'b1_without_bb' -mode = 'b1_with_bb' -#mode = 'b1_with_bb_legacy_macros' -#mode = 'b4_without_bb' -#mode = 'b4_from_b2_without_bb' -mode = 'b4_from_b2_with_bb' - -# Tolarances for checks [ip1, ip2, ip5, ip8] -tol_beta = [1e-3, 10e-2, 1e-3, 1e-2] -tol_sep = [1e-6, 1e-6, 1e-6, 1e-6] - -pm.make_links(force=True, links_dict={ - 'tracking_tools': '/afs/cern.ch/eng/tracking-tools', - 'modules': 'tracking_tools/modules', - 'tools': 'tracking_tools/tools', - 'beambeam_macros': 'tracking_tools/beambeam_macros', - 'errors': 'tracking_tools/errors'}) - -optics_file = 'opticsfile.29' - -check_betas_at_ips = True -check_separations_at_ips = True -save_intermediate_twiss = True - -# Check and load parameters -pm.checks_on_parameter_dict(mask_parameters) - -# Define configuration -(beam_to_configure, sequences_to_check, sequence_to_track, generate_b4_from_b2, - track_from_b4_mad_instance, enable_bb_python, enable_bb_legacy, - force_disable_check_separations_at_ips, - ) = pm.get_pymask_configuration(mode) -if force_disable_check_separations_at_ips: - check_separations_at_ips = False - -# Start mad -mad = Madx() - -# Build sequence -ost.build_sequence(mad, beam=beam_to_configure) - -# Apply optics -ost.apply_optics(mad, optics_file=optics_file) - -# Force disable beam-beam when needed -if not(enable_bb_legacy) and not(enable_bb_python): - mask_parameters['par_on_bb_switch'] = 0. - -# Pass parameters to mad -mad.set_variables_from_dict(params=mask_parameters) - -# Prepare sequences and attach beam -mad.call("modules/submodule_01a_preparation.madx") -mad.call("modules/submodule_01b_beam.madx") - -# Test machine before any change -twiss_dfs, other_data = ost.twiss_and_check(mad, sequences_to_check, - tol_beta=tol_beta, tol_sep=tol_sep, - twiss_fname='twiss_from_optics', - save_twiss_files=save_intermediate_twiss, - check_betas_at_ips=check_betas_at_ips, - check_separations_at_ips=check_separations_at_ips) - -# Set phase, apply and save crossing -mad.call("modules/submodule_01c_phase.madx") - -# Set optics-specific knobs -ost.set_optics_specific_knobs(mad, mode) - -# Crossing-save and some reference measurements -mad.input('exec, crossing_save') -mad.call("modules/submodule_01e_final.madx") - -# Test flat machine -mad.input('exec, crossing_disable') -twiss_dfs, other_data = ost.twiss_and_check(mad, sequences_to_check, - tol_beta=tol_beta, tol_sep=tol_sep, - twiss_fname='twiss_no_crossing', - save_twiss_files=save_intermediate_twiss, - check_betas_at_ips=check_betas_at_ips, check_separations_at_ips=check_separations_at_ips) -# Check flatness -flat_tol = 1e-6 -for ss in twiss_dfs.keys(): - tt = twiss_dfs[ss] - assert np.max(np.abs(tt.x)) < flat_tol - assert np.max(np.abs(tt.y)) < flat_tol - -# Check machine after crossing restore -mad.input('exec, crossing_restore') -twiss_dfs, other_data = ost.twiss_and_check(mad, sequences_to_check, - tol_beta=tol_beta, tol_sep=tol_sep, - twiss_fname='twiss_with_crossing', - save_twiss_files=save_intermediate_twiss, - check_betas_at_ips=check_betas_at_ips, check_separations_at_ips=check_separations_at_ips) - -mad.use(f'lhcb{beam_to_configure}') - -# # Call leveling module -# if mode=='b4_without_bb': -# print('Leveling not working in this mode!') -# else: -# mad.call("modules/module_02_lumilevel.madx") - - -mad.input('on_disp = 0') - -# Prepare bb dataframes -if enable_bb_python: - bb_dfs = pm.generate_bb_dataframes(mad, - ip_names=['ip1', 'ip2', 'ip5', 'ip8'], - harmonic_number=35640, - numberOfLRPerIRSide=[25, 20, 25, 20], - bunch_spacing_buckets=10, - numberOfHOSlices=11, - bunch_population_ppb=None, - sigmaz_m=None, - z_crab_twiss = 0, - remove_dummy_lenses=True) - - # Here the datafremes can be edited, e.g. to set bbb intensity - -# Generate mad instance for b4 -if generate_b4_from_b2: - mad_b4 = Madx() - ost.build_sequence(mad_b4, beam=4) - ost.apply_optics(mad_b4, optics_file=optics_file) - - pm.configure_b4_from_b2(mad_b4, mad) - - twiss_dfs_b2, other_data_b2 = ost.twiss_and_check(mad, - sequences_to_check=['lhcb2'], - tol_beta=tol_beta, tol_sep=tol_sep, - twiss_fname='twiss_b2_for_b4check', - save_twiss_files=save_intermediate_twiss, - check_betas_at_ips=check_betas_at_ips, check_separations_at_ips=False) - - twiss_dfs_b4, other_data_b4 = ost.twiss_and_check(mad_b4, - sequences_to_check=['lhcb2'], - tol_beta=tol_beta, tol_sep=tol_sep, - twiss_fname='twiss_b4_for_b4check', - save_twiss_files=save_intermediate_twiss, - check_betas_at_ips=check_betas_at_ips, check_separations_at_ips=False) - - - -# We working exclusively on the sequence to track -# Select mad object -if track_from_b4_mad_instance: - mad_track = mad_b4 -else: - mad_track = mad - -mad_collider = mad -del(mad) - -# Twiss machine to track -twiss_dfs, other_data = ost.twiss_and_check(mad_track, sequences_to_check, - tol_beta=tol_beta, tol_sep=tol_sep, - twiss_fname='twiss_track_intermediate', - save_twiss_files=save_intermediate_twiss, - check_betas_at_ips=check_betas_at_ips, check_separations_at_ips=False) - - -# Install bb lenses -if enable_bb_python: - if track_from_b4_mad_instance: - bb_df_track = bb_dfs['b4'] - assert(sequence_to_track=='lhcb2') - else: - bb_df_track = bb_dfs['b1'] - assert(sequence_to_track=='lhcb1') - - pm.install_lenses_in_sequence(mad_track, bb_df_track, sequence_to_track) - - # Disable bb - mad_track.globals.on_bb_charge = 0 -else: - bb_df_track = None - - -# Legacy bb macros -if enable_bb_legacy: - assert(beam_to_configure == 1) - assert(not(track_from_b4_mad_instance)) - assert(not(enable_bb_python)) - mad_track.call("modules/module_03_beambeam.madx") - -# # Install crab cavities -# mad_track.call("tools/enable_crabcavities.madx") - -# # Install crab cavities -#mad_track.call("modules/submodule_04_1a_install_crabs.madx") - -# Save references (orbit at IPs) -mad_track.call('modules/submodule_04_1b_save_references.madx') - -# Switch off dipersion correction knob -mad_track.globals.on_disp = 0. - -# Final use -mad_track.use(sequence_to_track) -# Disable use -mad_track._use = mad_track.use -mad_track.use = None - -# # Install and correct errors -# mad_track.call("modules/module_04_errors.madx") - -# Machine tuning (enables bb) -mad_track.call("modules/module_05_tuning.madx") - -# # Switch on crab cavities -# mad_track.globals.on_crab1 = mad_track.globals.par_crab1 -# mad_track.globals.on_crab5 = mad_track.globals.par_crab5 - -# Generate sixtrack -if enable_bb_legacy: - mad_track.call("modules/module_06_generate.madx") -else: - pm.generate_sixtrack_input(mad_track, - seq_name=sequence_to_track, - bb_df=bb_df_track, - output_folder='./', - reference_bunch_charge_sixtrack_ppb=( - mad_track.sequence[sequence_to_track].beam.npart), - emitnx_sixtrack_um=( - mad_track.sequence[sequence_to_track].beam.exn), - emitny_sixtrack_um=( - mad_track.sequence[sequence_to_track].beam.eyn), - sigz_sixtrack_m=( - mad_track.sequence[sequence_to_track].beam.sigt), - sige_sixtrack=( - mad_track.sequence[sequence_to_track].beam.sige), - ibeco_sixtrack=1, - ibtyp_sixtrack=0, - lhc_sixtrack=2, - ibbc_sixtrack=0, - radius_sixtrack_multip_conversion_mad=0.017, - skip_mad_use=True) - -# Get optics and orbit at start ring -optics_orbit_start_ring = pm.get_optics_and_orbit_at_start_ring( - mad_track, sequence_to_track, skip_mad_use=True) -with open('./optics_orbit_at_start_ring.pkl', 'wb') as fid: - pickle.dump(optics_orbit_start_ring, fid) - -# Generate pysixtrack lines -if enable_bb_legacy: - print('Pysixtrack line is not generated with bb legacy macros') -else: - pysix_fol_name = "./pysixtrack" - dct_pysxt = pm.generate_pysixtrack_line_with_bb(mad_track, - sequence_to_track, bb_df_track, - closed_orbit_method='from_mad', - pickle_lines_in_folder=pysix_fol_name, - skip_mad_use=True) diff --git a/python_examples/run3_collisions_python/old/knob_parameters.py b/python_examples/run3_collisions_python/old/knob_parameters.py deleted file mode 100644 index f2d6c02..0000000 --- a/python_examples/run3_collisions_python/old/knob_parameters.py +++ /dev/null @@ -1,29 +0,0 @@ - -knob_parameters = { - - #IP specific orbit settings - 'par_x1' : 150., - 'par_x5' : 150., - 'par_sep1' : 0, - 'par_sep5' : 0, - 'par_sep2h' : 1, - 'par_sep2v' : 0, - 'par_x2h' : 0, - 'par_x2v' : 200, - 'par_sep8h' : 0, - 'par_sep8v' : 1, - 'par_x8h' : 0, - 'par_x8v' : 135, - - # Dispersion correction knob - 'par_on_disp' : 1, # Could be optics-dependent - - # Magnets of the experiments - 'par_on_alice' : 1, - 'par_on_lhcb' : 1, - - 'par_on_sol_atlas' : 0, - 'par_on_sol_cms' : 0, - 'par_on_sol_alice' : 0, - -} diff --git a/python_examples/run3_collisions_python/old/mask_parameters.py b/python_examples/run3_collisions_python/old/mask_parameters.py deleted file mode 100644 index 017ce09..0000000 --- a/python_examples/run3_collisions_python/old/mask_parameters.py +++ /dev/null @@ -1,71 +0,0 @@ - -mask_parameters = { - 'par_verbose' : 1, - - # Beam parameters - 'par_beam_norm_emit' : 2.5, # [um] - 'par_beam_sigt' : 0.075, # [m] - 'par_beam_sige' : 1.1e-4, # [-] - 'par_beam_npart' : 1.16e11, # [-] - 'par_beam_energy_tot' : 7000, # [GeV] - - # Settings - 'par_oct_current' : 350., # [A] - 'par_chromaticity' : 15., # [-] (Q':5 for colliding bunches, Q':15 for non-colliding bunches) - 'par_vrf_total' : 16., # [MV] - - # Tunes - 'par_qx0' : 62.31, - 'par_qy0' : 60.32, - - - #*************************# - # Beam-beam configuration # - #*************************# - - 'par_on_bb_switch' : 1, - 'par_match_with_bb' : 0, # should be off at collision - 'par_b_t_dist' : 25., # bunch separation [ns] - 'par_n_inside_D1' : 5, # default value for the number of additionnal parasitic encounters inside D1 - - 'par_nho_IR1' : 11, # number of slices for head-on in IR1 (between 0 and 201) - 'par_nho_IR2' : 11, # number of slices for head-on in IR2 (between 0 and 201) - 'par_nho_IR5' : 11, # number of slices for head-on in IR5 (between 0 and 201) - 'par_nho_IR8' : 11, # number of slices for head-on in IR8 (between 0 and 201) - - #*************************# - # Leveling in IP8 # - #*************************# - - # This variables set the leveled luminosity in IP8 (considered if par_on_collision:1) - 'par_lumi_ip8' : 2e33, #[Hz/cm2] - - # These variables define the number of Head-On collisions in the 4 IPs - 'par_nco_IP1' : 2592, - 'par_nco_IP2' : 2288, - 'par_nco_IP5' : 2592, - 'par_nco_IP8' : 2396, - - #*************************# - # Errors and corrections # - #*************************# - - # Select seed for errors - 'par_myseed' : 0, - - # Set this flag to correct the errors of D2 in the NLC (warning: for now only correcting b3 of D2, still in development) - 'par_correct_for_D2' : 0, - # Set this flag to correct the errors of MCBXF in the NLC (warning: this might be less reproducable in reality, use with care) - 'par_correct_for_MCBX' : 0, - - 'par_on_errors_LHC' : 0, - 'par_on_errors_MBH' : 0, - 'par_on_errors_Q5' : 0, - 'par_on_errors_Q4' : 0, - 'par_on_errors_D2' : 0, - 'par_on_errors_D1' : 0, - 'par_on_errors_IT' : 0, - 'par_on_errors_MCBRD' : 0, - 'par_on_errors_MCBXF' : 0, - -} diff --git a/python_examples/run3_collisions_python/old/optics_specific_tools.py b/python_examples/run3_collisions_python/old/optics_specific_tools.py deleted file mode 100644 index b990ad2..0000000 --- a/python_examples/run3_collisions_python/old/optics_specific_tools.py +++ /dev/null @@ -1,186 +0,0 @@ -import pymask as pm - -# The parts marked by (*) in following need to be -# adapted according to knob definitions - -def build_sequence(mad, beam): - - slicefactor = 2 - - pm.make_links(force=True, links_dict={ - 'optics_indep_macros.madx': 'tools/optics_indep_macros.madx', - 'macro.madx': ('/afs/cern.ch/user/s/sterbini/public/' - 'tracking_tools/tools/macro.madx'), - 'optics_runII': '/afs/cern.ch/eng/lhc/optics/runII', - 'optics_runIII': '/afs/cern.ch/eng/lhc/optics/runIII',}) - - - mylhcbeam = int(beam) - - mad.input('ver_lhc_run = 3') - - mad.input(f'mylhcbeam = {beam}') - mad.input('option, -echo,warn, -info;') - - # optics dependent macros - mad.call('macro.madx') - # optics independent macros - mad.call('optics_indep_macros.madx') - - assert mylhcbeam in [1, 2, 4], "Invalid mylhcbeam (it should be in [1, 2, 4])" - - if mylhcbeam in [1, 2]: - mad.call('optics_runII/2018/lhc_as-built.seq') - else: - mad.call('optics_runII/2018/lhcb4_as-built.seq') - - # New IR7 MQW layout and cabling - mad.call('optics_runIII/RunIII_dev/IR7-Run3seqedit.madx') - - # Makethin part - if slicefactor > 0: - # the variable in the macro is slicefactor - mad.input(f'slicefactor={slicefactor};') - mad.call('optics_runII/2018/toolkit/myslice.madx') - mad.beam() - for my_sequence in ['lhcb1','lhcb2']: - if my_sequence in list(mad.sequence): - mad.input(f'use, sequence={my_sequence}; makethin, sequence={my_sequence}, style=teapot, makedipedge=false;') - else: - warnings.warn('The sequences are not thin!') - - # Cycling w.r.t. to IP3 (mandatory to find closed orbit in collision in the presence of errors) - for my_sequence in ['lhcb1','lhcb2']: - if my_sequence in list(mad.sequence): - mad.input(f'seqedit, sequence={my_sequence}; flatten; cycle, start=IP3; flatten; endedit;') - -def apply_optics(mad, optics_file): - pm.make_links(force=True, links_dict={ - 'optics.madx' : 'optics_runIII/RunIII_dev/2022_V1/PROTON/' + optics_file}) - mad.call('optics.madx') - -def check_beta_at_ips_against_madvars(beam, twiss_df, variable_dicts, tol): - twiss_value_checks=[] - for iip, ip in enumerate([1,2,5,8]): - for plane in ['x', 'y']: - # (*) Adapet based on knob definitions - twiss_value_checks.append({ - 'element_name': f'ip{ip}:1', - 'keyword': f'bet{plane}', - 'varname': f'bet{plane}ip{ip}b{beam}', - 'tol': tol[iip]}) - - pm.check_twiss_against_madvars(twiss_value_checks, twiss_df, variable_dicts) - - -def set_optics_specific_knobs(mad, mode=None): - - from knob_parameters import knob_parameters as kp - - mad.set_variables_from_dict(params=kp) - - # Set IP knobs - mad.globals['on_x1'] = kp['par_x1'] - mad.globals['on_sep1'] = kp['par_sep1'] - - mad.globals['on_x2h'] = kp['par_x2h'] - mad.globals['on_x2v'] = kp['par_x2v'] - mad.globals['on_sep2h'] = kp['par_sep2h'] - mad.globals['on_sep2v'] = kp['par_sep2v'] - - mad.globals['on_x5'] = kp['par_x5'] - mad.globals['on_sep5'] = kp['par_sep5'] - - mad.globals['on_x8h'] = kp['par_x8h'] - mad.globals['on_x8v'] = kp['par_x8v'] - mad.globals['on_sep8h'] = kp['par_sep8h'] - mad.globals['on_sep8v'] = kp['par_sep8v'] - - mad.globals['on_disp'] = kp['par_on_disp'] - - # A check - if mad.globals.nrj < 500: - assert kp['par_on_disp'] == 0 - - # Spectrometers at experiments - if kp['par_on_alice'] == 1: - mad.globals.on_alice = 7000./mad.globals.nrj - if kp['par_on_lhcb'] == 1: - mad.globals.on_lhcb = 7000./mad.globals.nrj - - # Solenoids at experiments - mad.globals.on_sol_atlas = kp['par_on_sol_atlas'] - mad.globals.on_sol_cms = kp['par_on_sol_cms'] - mad.globals.on_sol_alice = kp['par_on_sol_alice'] - -def check_separations_at_ips_against_madvars(twiss_df_b1, twiss_df_b2, - variables_dict, tol): - - separations_to_check = [] - for iip, ip in enumerate([2,8]): - for plane in ['x', 'y']: - # (*) Adapet based on knob definitions - separations_to_check.append({ - 'element_name': f'ip{ip}:1', - 'scale_factor': -2*1e-3, - 'plane': plane, - # knobs like on_sep1h, onsep8v etc - 'varname': f'on_sep{ip}'+{'x':'h', 'y':'v'}[plane], - 'tol': tol[iip]}) - separations_to_check.append({ # IP1 - 'element_name': f'ip1:1', - 'scale_factor': -2*1e-3, - 'plane': 'x', - 'varname': 'on_sep1', - 'tol': tol[0]}) - separations_to_check.append({ # IP5 - 'element_name': f'ip5:1', - 'scale_factor': -2*1e-3, - 'plane': 'y', - 'varname': 'on_sep5', - 'tol': tol[2]}) - pm.check_separations_against_madvars(separations_to_check, - twiss_df_b1, twiss_df_b2, variables_dict) - -def twiss_and_check(mad, sequences_to_check, twiss_fname, - tol_beta=1e-3, tol_sep=1e-6, save_twiss_files=True, - check_betas_at_ips=True, check_separations_at_ips=True): - - var_dict = mad.get_variables_dicts() - twiss_dfs = {} - summ_dfs = {} - for ss in sequences_to_check: - mad.use(ss) - mad.twiss() - tdf = mad.get_twiss_df('twiss') - twiss_dfs[ss] = tdf - sdf = mad.get_summ_df('summ') - summ_dfs[ss] = sdf - - if save_twiss_files: - for ss in sequences_to_check: - tt = twiss_dfs[ss] - if twiss_fname is not None: - tt.to_parquet(twiss_fname + f'_seq_{ss}.parquet') - - if check_betas_at_ips: - for ss in sequences_to_check: - tt = twiss_dfs[ss] - check_beta_at_ips_against_madvars(beam=ss[-1], - twiss_df=tt, - variable_dicts=var_dict, - tol=tol_beta) - print('IP beta test against knobs passed!') - - if check_separations_at_ips: - twiss_df_b1 = twiss_dfs['lhcb1'] - twiss_df_b2 = twiss_dfs['lhcb2'] - check_separations_at_ips_against_madvars(twiss_df_b1, twiss_df_b2, - var_dict, tol=tol_sep) - print('IP separation test against knobs passed!') - - other_data = {} - other_data.update(var_dict) - other_data['summ_dfs'] = summ_dfs - - return twiss_dfs, other_data diff --git a/python_examples/run3_collisions_python/old/t000_plots_from_twiss.py b/python_examples/run3_collisions_python/old/t000_plots_from_twiss.py deleted file mode 100644 index 50c9863..0000000 --- a/python_examples/run3_collisions_python/old/t000_plots_from_twiss.py +++ /dev/null @@ -1,38 +0,0 @@ -import numpy as np -import pandas as pd -import matplotlib.pyplot as plt - -twiss_df = pd.read_parquet('./twiss_with_crossing_seq_lhcb1.parquet') - - -plt.close('all') - -# %% -fig = plt.figure(1, figsize=(6.4*1.6, 4.8*1.5)) -axbetx = fig.add_subplot(2,2,1) -axbetx.plot(twiss_df['s'], twiss_df['betx']) -axbetx.set_xlabel('s') -axbetx.set_ylabel(r'$\beta_x$ [m]') -axbetx.ticklabel_format(style='sci', scilimits=(0, 0), axis='y') - -# %% -axbety = fig.add_subplot(2,2,2, sharex=axbetx, sharey=axbetx) -axbety.plot(twiss_df['s'], twiss_df['bety']) -axbety.set_xlabel('s') -axbety.set_ylabel(r'$\beta_y$ [m]') -axbety.ticklabel_format(style='sci', scilimits=(0, 0), axis='y') -axbety.legend() -# %% - -axcox = fig.add_subplot(2,2,3, sharex=axbetx) -axcox.plot(twiss_df['s'], twiss_df['x']) -axcox.set_xlabel('s') -axcox.set_ylabel('x [m]') -#plt.xlim(5500,7500) -# %% -axcoy = fig.add_subplot(2,2,4, sharex=axbetx, sharey=axcox) -axcoy.plot(twiss_df['s'], twiss_df['y']) -axcoy.set_xlabel('s') -axcoy.set_ylabel('y [m]') - -plt.show() diff --git a/python_examples/run3_collisions_python/old/t001_check_b4_against_b2.py b/python_examples/run3_collisions_python/old/t001_check_b4_against_b2.py deleted file mode 100644 index d8aa29d..0000000 --- a/python_examples/run3_collisions_python/old/t001_check_b4_against_b2.py +++ /dev/null @@ -1,44 +0,0 @@ -import numpy as np -import pandas as pd -import matplotlib.pyplot as plt - -twiss_df_b2 = pd.read_parquet('./twiss_b2_for_b4check_seq_lhcb2.parquet') -twiss_df_b4 = pd.read_parquet('./twiss_b4_for_b4check_seq_lhcb2.parquet') - - -plt.close('all') - -# %% -fig = plt.figure(1, figsize=(6.4*1.6, 4.8*1.5)) -axbetx = fig.add_subplot(2,2,1) -axbetx.plot(twiss_df_b4['s'][-1]-twiss_df_b4['s'],twiss_df_b4['betx'],'b') -axbetx.plot(twiss_df_b2['s'], twiss_df_b2['betx'],'--r') -axbetx.set_xlabel('s') -axbetx.set_ylabel(r'$\beta_x$ [m]') -axbetx.ticklabel_format(style='sci', scilimits=(0, 0), axis='y') - -# %% -axbety = fig.add_subplot(2,2,2, sharex=axbetx, sharey=axbetx) -axbety.plot(twiss_df_b4['s'][-1]-twiss_df_b4['s'],twiss_df_b4['bety'],'b', - label='b4') -axbety.plot(twiss_df_b2['s'], twiss_df_b2['bety'],'--r', label='b2') -axbety.set_xlabel('s') -axbety.set_ylabel(r'$\beta_y$ [m]') -axbety.ticklabel_format(style='sci', scilimits=(0, 0), axis='y') -axbety.legend() -# %% - -axcox = fig.add_subplot(2,2,3, sharex=axbetx) -axcox.plot(twiss_df_b4['s'][-1]-twiss_df_b4['s'],-twiss_df_b4['x'],'b') -axcox.plot(twiss_df_b2['s'], twiss_df_b2['x'], '--r') -axcox.set_xlabel('s') -axcox.set_ylabel('x [m]') -#plt.xlim(5500,7500) -# %% -axcoy = fig.add_subplot(2,2,4, sharex=axbetx, sharey=axcox) -axcoy.plot(twiss_df_b4['s'][-1]-twiss_df_b4['s'],twiss_df_b4['y'],'b') -axcoy.plot(twiss_df_b2['s'], twiss_df_b2['y'], '--r') -axcoy.set_xlabel('s') -axcoy.set_ylabel('y [m]') - -plt.show() diff --git a/python_examples/run3_collisions_python/old/t003_fc_to_fort.py b/python_examples/run3_collisions_python/old/t003_fc_to_fort.py deleted file mode 100644 index 5c77e19..0000000 --- a/python_examples/run3_collisions_python/old/t003_fc_to_fort.py +++ /dev/null @@ -1,25 +0,0 @@ -import os -import shutil - -folder_list = [ - './', - '../hl_lhc_collision/', - '../hl_lhc_collision/reference', - '../hl_lhc_collision_nobb_b4'] - -for sixtrack_input_folder in folder_list: - for iff in [2,8,16,34]: - os.system(f"rm {sixtrack_input_folder}/fort.{iff}") - try: - shutil.copy(sixtrack_input_folder + f"/fc.{iff}", - sixtrack_input_folder + f"/fort.{iff}") - except Exception: - print(f"/fc.{iff} not found!") - - with open(sixtrack_input_folder + "/fort.3", "w") as fout: - with open("fort_parts/fort_beginning.3", "r") as fid_fort3b: - fout.write(fid_fort3b.read()) - with open(sixtrack_input_folder + "/fc.3", "r") as fid_fc3: - fout.write(fid_fc3.read()) - with open("fort_parts/fort_end.3", "r") as fid_fort3e: - fout.write(fid_fort3e.read())