Merge branch 'release_v1.0.0'

Author: sterbini

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
+    bb_df['yma'] = yma

pymask/pymasktools.py

@@ -182,7 +182,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 +194,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 +234,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 +260,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 +298,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}",
@@ -370,24 +378,24 @@ def get_optics_and_orbit_at_start_ring(mad, seq_name, with_bb_forces=False,
             }
     return optics_at_start_ring
 
-def generate_pysixtrack_line_with_bb(mad, seq_name, bb_df,
+def generate_xline_with_bb(mad, seq_name, bb_df,
         closed_orbit_method='from_mad', pickle_lines_in_folder=None,
         skip_mad_use=False):
 
     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])
+    # Build xline model
+    import xline
+    line = xline.Line.from_madx_sequence(
+        mad.sequence[seq_name], apply_madx_errors=True)
 
     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(
+            xline.elements.Cavity)
     for cc, nn in zip(cavities, cav_names):
         if cc.frequency ==0.:
             ii_mad = mad.sequence[seq_name].element_names().index(nn)
@@ -397,37 +405,42 @@ def generate_pysixtrack_line_with_bb(mad, seq_name, bb_df,
 
     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(
+    line.disable_beambeam()
+    part_on_CO = 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()
+    line.enable_beambeam()
 
-    pysxt_line_bb_dipole_cancelled = pysxt_line.copy()
+    line_bb_dipole_cancelled = line.copy()
 
-    pysxt_line_bb_dipole_cancelled.beambeam_store_closed_orbit_and_dipolar_kicks(
+    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,
+    xline_dict = {
+            'line_bb_dipole_not_cancelled': line,
+            'line_bb_dipole_cancelled': 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
-
-
+        xline_fol_name = pickle_lines_in_folder
+        os.makedirs(xline_fol_name, exist_ok=True)
+        
+        line.to_json(xline_fol_name + "/line_bb_dipole_not_cancelled.json", keepextra=True)
+        line_bb_dipole_cancelled.to_json(xline_fol_name + "/line_bb_dipole_cancelled.json", keepextra=True)
+        part_on_CO.to_json(xline_fol_name + "/particle_on_closed_orbit.json")
+
+    return xline_dict
+
+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")

python_examples/clean_all

@@ -9,7 +9,7 @@ rm twiss_bb
 rm last_twiss.*.gz
 rm fort.*
 rm *.tfs
-rm -r pysixtrack
+rm -r xline 
 rm -r temp
 rm -r __pycache__