Skip to content

Commit 1790cea

Browse files
sterbinisterbini
committed
Same 000_pymask.py for ions and protons
1 parent 84e4f4b commit 1790cea

File tree

3 files changed

+111
-52
lines changed

3 files changed

+111
-52
lines changed

‎python_examples/hl_lhc_collisions_python/000_pymask.py‎

Lines changed: 39 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,21 @@
8787

8888
# Attach beam to sequences
8989
mad.globals.nrj = configuration['beam_energy_tot']
90-
gamma_rel = configuration['beam_energy_tot']/mad.globals.pmass
90+
particle_type = 'proton'
91+
92+
if 'beam_mass' in configuration.keys():
93+
beam_mass = configuration['beam_mass']
94+
particle_type = 'ion'
95+
else:
96+
beam_mass = mad.globals.pmass # proton mass
97+
98+
if 'beam_charge' in configuration.keys():
99+
beam_charge = configuration['beam_charge']
100+
particle_type = 'ion'
101+
else:
102+
beam_charge = 1.
103+
104+
gamma_rel = (beam_charge*configuration['beam_energy_tot'])/beam_mass
91105
for ss in mad.sequence.keys():
92106
# bv and bv_aux flags
93107
if ss == 'lhcb1':
@@ -100,14 +114,16 @@
100114

101115
mad.globals['bv_aux'] = ss_bv_aux
102116
mad.input(f'''
103-
beam, particle=proton,sequence={ss},
104-
energy={configuration['beam_energy_tot']},
117+
beam, particle={particle_type},sequence={ss},
118+
energy={configuration['beam_energy_tot']*beam_charge},
105119
sigt={configuration['beam_sigt']},
106120
bv={ss_beam_bv},
107121
npart={configuration['beam_npart']},
108122
sige={configuration['beam_sige']},
109123
ex={configuration['beam_norm_emit_x'] * 1e-6 / gamma_rel},
110124
ey={configuration['beam_norm_emit_y'] * 1e-6 / gamma_rel},
125+
mass={beam_mass},
126+
charge={beam_charge},
111127
''')
112128

113129

@@ -177,6 +193,8 @@
177193
if mode=='b4_without_bb':
178194
print('Leveling not working in this mode!')
179195
else:
196+
if particle_type == 'ion': # the legacy macro for BB have been checked but not maintained
197+
raise ValueError
180198
# Luminosity levelling
181199
vars_for_legacy_level = ['lumi_ip8',
182200
'nco_IP1', 'nco_IP2', 'nco_IP5', 'nco_IP8']
@@ -267,6 +285,13 @@
267285
save_twiss_files=save_intermediate_twiss,
268286
check_betas_at_ips=check_betas_at_ips, check_separations_at_ips=False)
269287

288+
# For B1, to be generalized for B4
289+
if 'filling_scheme_json' in configuration.keys():
290+
filling_scheme_json = configuration['filling_scheme_json']
291+
bunch_to_track = configuration['bunch_to_track']
292+
bb_schedule_to_track_b1 = ost.create_bb_shedule_to_track(
293+
filling_scheme_json,bunch_to_track, beam='1')
294+
bb_dfs['b1']=ost.filter_bb_df(bb_dfs['b1'],bb_schedule_to_track_b1)
270295

271296
##################################################
272297
# Select mad instance for tracking configuration #
@@ -508,3 +533,14 @@
508533
closed_orbit_method='from_mad',
509534
pickle_lines_in_folder=pysix_fol_name,
510535
skip_mad_use=True)
536+
537+
#############################
538+
# Save final twiss #
539+
#############################
540+
541+
mad_track.globals.on_bb_charge = 0
542+
mad_track.twiss()
543+
tdf = mad_track.get_twiss_df('twiss')
544+
sdf = mad_track.get_summ_df('summ')
545+
tdf.to_parquet('final_twiss_BBOFF.parquet')
546+
sdf.to_parquet('final_summ_BBOFF.parquet')

‎python_examples/ions_python/000_pymask.py‎

Lines changed: 33 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,21 @@
8787

8888
# Attach beam to sequences
8989
mad.globals.nrj = configuration['beam_energy_tot']
90-
gamma_rel = (configuration['beam_charge']*configuration['beam_energy_tot'])/configuration['beam_mass']
90+
particle_type = 'proton'
91+
92+
if 'beam_mass' in configuration.keys():
93+
beam_mass = configuration['beam_mass']
94+
particle_type = 'ion'
95+
else:
96+
beam_mass = mad.globals.pmass # proton mass
97+
98+
if 'beam_charge' in configuration.keys():
99+
beam_charge = configuration['beam_charge']
100+
particle_type = 'ion'
101+
else:
102+
beam_charge = 1.
103+
104+
gamma_rel = (beam_charge*configuration['beam_energy_tot'])/beam_mass
91105
for ss in mad.sequence.keys():
92106
# bv and bv_aux flags
93107
if ss == 'lhcb1':
@@ -100,16 +114,16 @@
100114

101115
mad.globals['bv_aux'] = ss_bv_aux
102116
mad.input(f'''
103-
beam, particle=ion,sequence={ss},
104-
energy={configuration['beam_energy_tot']*configuration['beam_charge']},
117+
beam, particle={particle_type},sequence={ss},
118+
energy={configuration['beam_energy_tot']*beam_charge},
105119
sigt={configuration['beam_sigt']},
106120
bv={ss_beam_bv},
107121
npart={configuration['beam_npart']},
108122
sige={configuration['beam_sige']},
109123
ex={configuration['beam_norm_emit_x'] * 1e-6 / gamma_rel},
110124
ey={configuration['beam_norm_emit_y'] * 1e-6 / gamma_rel},
111-
mass = {configuration['beam_mass']},
112-
charge={configuration['beam_charge']},
125+
mass={beam_mass},
126+
charge={beam_charge},
113127
''')
114128

115129

@@ -179,45 +193,14 @@
179193
if mode=='b4_without_bb':
180194
print('Leveling not working in this mode!')
181195
else:
182-
# Modified module 2 for offset levelling in IP1,2,5,8
183-
vars_for_legacy_level = ['lumi_ip8','lumi_ip5', 'lumi_ip1', 'lumi_ip2',
196+
if particle_type == 'ion': # the legacy macro for BB have been checked but not maintained
197+
raise ValueError
198+
# Luminosity levelling
199+
vars_for_legacy_level = ['lumi_ip8',
184200
'nco_IP1', 'nco_IP2', 'nco_IP5', 'nco_IP8']
185201
mad.set_variables_from_dict({
186202
'par_'+kk: configuration[kk] for kk in vars_for_legacy_level})
187-
mad.input('''
188-
print, text="";
189-
print, text="";
190-
print, text="+++++++++++++++++++++++++++++++";
191-
print, text="++ START MODULE 2: LEVELLING ++";
192-
print, text="+++++++++++++++++++++++++++++++";
193-
print, text="";
194-
print, text="";
195-
196-
call, file="beambeam_macros/macro_bb.madx"; ! macros for beam-beam
197-
198-
exec, DEFINE_BB_PARAM; !Define main beam-beam parameters
199-
200-
!Switch on Xscheme in precollision
201-
on_disp:=0;
202-
halo1=0;halo2=0;halo5=0;halo8=0; !halo collision at 5 sigma's in Alice
203-
!number of collision/turn at IP1/2/5/8
204-
nco_IP1 = par_nco_IP1;
205-
nco_IP5 = par_nco_IP5;
206-
nco_IP2 = par_nco_IP2;
207-
nco_IP8 = par_nco_IP8;
208-
exec, LEVEL_PARALLEL_OFFSET_FOR(par_lumi_ip8, 8); value,halo8;
209-
exec, LEVEL_PARALLEL_OFFSET_FOR(par_lumi_ip1, 1); value,halo1;
210-
exec, LEVEL_PARALLEL_OFFSET_FOR(par_lumi_ip2, 2); value,halo2;
211-
exec, LEVEL_PARALLEL_OFFSET_FOR(par_lumi_ip5, 5); value,halo5;
212-
!Redefine the on_sep's accordingly
213-
exec, CALCULATE_XSCHEME(halo1,halo2,halo5,halo8);
214-
! Saving new crossing scheme with separation
215-
on_disp=on_dispaux; ! reset on_disp before saving
216-
exec, crossing_save;
217-
218-
if (mylhcbeam==1) { use, sequence=lhcb1; } else { use, sequence=lhcb2; };
219-
220-
''')
203+
mad.input("call, file='modules/module_02_lumilevel.madx';")
221204
else:
222205
print('Start pythonic leveling:')
223206
ost.lumi_control(mad, twiss_dfs, configuration, knob_names)
@@ -303,10 +286,12 @@
303286
check_betas_at_ips=check_betas_at_ips, check_separations_at_ips=False)
304287

305288
# For B1, to be generalized for B4
306-
filling_scheme_json = configuration['filling_scheme_json']
307-
bunch_to_track = configuration['bunch_to_track']
308-
bb_schedule_to_track_b1 = ost.create_bb_shedule_to_track(filling_scheme_json,bunch_to_track, beam='1')
309-
bb_dfs['b1']=ost.filter_bb_df(bb_dfs['b1'],bb_schedule_to_track_b1)
289+
if 'filling_scheme_json' in configuration.keys():
290+
filling_scheme_json = configuration['filling_scheme_json']
291+
bunch_to_track = configuration['bunch_to_track']
292+
bb_schedule_to_track_b1 = ost.create_bb_shedule_to_track(
293+
filling_scheme_json,bunch_to_track, beam='1')
294+
bb_dfs['b1']=ost.filter_bb_df(bb_dfs['b1'],bb_schedule_to_track_b1)
310295

311296
##################################################
312297
# Select mad instance for tracking configuration #
@@ -549,7 +534,9 @@
549534
pickle_lines_in_folder=pysix_fol_name,
550535
skip_mad_use=True)
551536

552-
##### Save final twiss
537+
#############################
538+
# Save final twiss #
539+
#############################
553540

554541
mad_track.globals.on_bb_charge = 0
555542
mad_track.twiss()

‎python_examples/run3_collisions_python/000_pymask.py‎

Lines changed: 39 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,21 @@
8787

8888
# Attach beam to sequences
8989
mad.globals.nrj = configuration['beam_energy_tot']
90-
gamma_rel = configuration['beam_energy_tot']/mad.globals.pmass
90+
particle_type = 'proton'
91+
92+
if 'beam_mass' in configuration.keys():
93+
beam_mass = configuration['beam_mass']
94+
particle_type = 'ion'
95+
else:
96+
beam_mass = mad.globals.pmass # proton mass
97+
98+
if 'beam_charge' in configuration.keys():
99+
beam_charge = configuration['beam_charge']
100+
particle_type = 'ion'
101+
else:
102+
beam_charge = 1.
103+
104+
gamma_rel = (beam_charge*configuration['beam_energy_tot'])/beam_mass
91105
for ss in mad.sequence.keys():
92106
# bv and bv_aux flags
93107
if ss == 'lhcb1':
@@ -100,14 +114,16 @@
100114

101115
mad.globals['bv_aux'] = ss_bv_aux
102116
mad.input(f'''
103-
beam, particle=proton,sequence={ss},
104-
energy={configuration['beam_energy_tot']},
117+
beam, particle={particle_type},sequence={ss},
118+
energy={configuration['beam_energy_tot']*beam_charge},
105119
sigt={configuration['beam_sigt']},
106120
bv={ss_beam_bv},
107121
npart={configuration['beam_npart']},
108122
sige={configuration['beam_sige']},
109123
ex={configuration['beam_norm_emit_x'] * 1e-6 / gamma_rel},
110124
ey={configuration['beam_norm_emit_y'] * 1e-6 / gamma_rel},
125+
mass={beam_mass},
126+
charge={beam_charge},
111127
''')
112128

113129

@@ -177,6 +193,8 @@
177193
if mode=='b4_without_bb':
178194
print('Leveling not working in this mode!')
179195
else:
196+
if particle_type == 'ion': # the legacy macro for BB have been checked but not maintained
197+
raise ValueError
180198
# Luminosity levelling
181199
vars_for_legacy_level = ['lumi_ip8',
182200
'nco_IP1', 'nco_IP2', 'nco_IP5', 'nco_IP8']
@@ -267,6 +285,13 @@
267285
save_twiss_files=save_intermediate_twiss,
268286
check_betas_at_ips=check_betas_at_ips, check_separations_at_ips=False)
269287

288+
# For B1, to be generalized for B4
289+
if 'filling_scheme_json' in configuration.keys():
290+
filling_scheme_json = configuration['filling_scheme_json']
291+
bunch_to_track = configuration['bunch_to_track']
292+
bb_schedule_to_track_b1 = ost.create_bb_shedule_to_track(
293+
filling_scheme_json,bunch_to_track, beam='1')
294+
bb_dfs['b1']=ost.filter_bb_df(bb_dfs['b1'],bb_schedule_to_track_b1)
270295

271296
##################################################
272297
# Select mad instance for tracking configuration #
@@ -508,3 +533,14 @@
508533
closed_orbit_method='from_mad',
509534
pickle_lines_in_folder=pysix_fol_name,
510535
skip_mad_use=True)
536+
537+
#############################
538+
# Save final twiss #
539+
#############################
540+
541+
mad_track.globals.on_bb_charge = 0
542+
mad_track.twiss()
543+
tdf = mad_track.get_twiss_df('twiss')
544+
sdf = mad_track.get_summ_df('summ')
545+
tdf.to_parquet('final_twiss_BBOFF.parquet')
546+
sdf.to_parquet('final_summ_BBOFF.parquet')

0 commit comments

Comments
 (0)