Merge pull request #82 from lhcopt/adapt_to_new_xfield_interface

sterbini · web-flow · commit ad9d09941eb9 · 2022-08-17T15:48:31.000+02:00
Adapt to new xfield interface
diff --git a/make_it.sh b/make_it.sh
@@ -1,14 +1,20 @@
 wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
 bash Miniconda3-latest-Linux-x86_64.sh -b  -p ./miniconda -f
 source miniconda/bin/activate
-python -m pip install ipython numpy scipy pandas psutil xsuite
+python -m pip install ipython numpy scipy pandas psutil xsuite matplotlib
 git clone -b release/v0.1.0 git@github.com:xsuite/tree_maker.git
 python -m pip install -e tree_maker
 git clone git@github.com:lhcopt/lhcmask.git -b release/v1.3.3
 python -m pip install -e lhcmask
+python -m pip install sixtracktools
+python -m pip install NAFFlib
 git clone git@github.com:lhcopt/lhcerrors.git
 git clone git@github.com:lhcopt/lhctoolkit.git
 git clone git@github.com:lhcopt/hllhc15.git
 git clone git@github.com:lhcopt/hllhc14.git
+# you need gfortran
+# `sudo yum install gcc-gfortran`
+# then `cd beambeam_macros`
+# and `gfortran headonslice.f -o headonslice`
 git clone git@github.com:lhcopt/beambeam_macros.git
 git clone $(whoami)@lxplus.cern.ch:/afs/cern.ch/eng/lhc/optics/runIII
diff --git a/pymask/beambeam.py b/pymask/beambeam.py
@@ -624,54 +624,40 @@ def setup_beam_beam_in_line(
 
     for ii, (ee, eename) in enumerate(zip(line.elements, line.element_names)):
         if isinstance(ee, xf.BeamBeamBiGaussian2D):
-            ee.n_particles=bb_df.loc[eename, 'other_num_particles']
-            ee.q0 = bb_df.loc[eename, 'other_particle_charge']
-            ee.sigma_x = np.sqrt(bb_df.loc[eename, 'other_Sigma_11'])
-            ee.sigma_y = np.sqrt(bb_df.loc[eename, 'other_Sigma_33'])
-            ee.beta0 = bb_df.loc[eename, 'other_relativistic_beta']
-            ee.mean_x = bb_df.loc[eename, 'separation_x']
-            ee.mean_y = bb_df.loc[eename, 'separation_y']
+            ee.other_beam_num_particles=bb_df.loc[eename, 'other_num_particles']
+            ee.other_beam_q0 = bb_df.loc[eename, 'other_particle_charge']
+            ee.other_beam_Sigma_11 = bb_df.loc[eename, 'other_Sigma_11']
+            ee.other_beam_Sigma_33 = bb_df.loc[eename, 'other_Sigma_33']
+            ee.other_beam_beta0 = bb_df.loc[eename, 'other_relativistic_beta']
+            ee.other_beam_shift_x = bb_df.loc[eename, 'separation_x']
+            ee.other_beam_shift_y = bb_df.loc[eename, 'separation_y']
         if isinstance(ee, xf.BeamBeamBiGaussian3D):
             params = {}
             params['phi'] = bb_df.loc[eename, 'phi']
             params['alpha'] =  bb_df.loc[eename, 'alpha']
-            params['x_bb_co'] =  bb_df.loc[eename, 'separation_x']
-            params['y_bb_co'] =  bb_df.loc[eename, 'separation_y']
-            # TODO update xtrack interface to separate charge and b. population
-            params['charge_slices'] =  [(bb_df.loc[eename, 'other_num_particles']
-                                 * bb_df.loc[eename, 'other_particle_charge'])]
-            params['zeta_slices'] =  [0.0]
-            params['sigma_11'] =  bb_df.loc[eename, 'other_Sigma_11']
-            params['sigma_12'] =  bb_df.loc[eename, 'other_Sigma_12']
-            params['sigma_13'] =  bb_df.loc[eename, 'other_Sigma_13']
-            params['sigma_14'] =  bb_df.loc[eename, 'other_Sigma_14']
-            params['sigma_22'] =  bb_df.loc[eename, 'other_Sigma_22']
-            params['sigma_23'] =  bb_df.loc[eename, 'other_Sigma_23']
-            params['sigma_24'] =  bb_df.loc[eename, 'other_Sigma_24']
-            params['sigma_33'] =  bb_df.loc[eename, 'other_Sigma_33']
-            params['sigma_34'] =  bb_df.loc[eename, 'other_Sigma_34']
-            params['sigma_44'] =  bb_df.loc[eename, 'other_Sigma_44']
+            params['other_beam_shift_x'] =  bb_df.loc[eename, 'separation_x']
+            params['other_beam_shift_y'] =  bb_df.loc[eename, 'separation_y']
+            params['slices_other_beam_num_particles'] =  [bb_df.loc[eename, 'other_num_particles']]
+            params['other_beam_q0'] =  bb_df.loc[eename, 'other_particle_charge']
+            params['slices_other_beam_zeta_center'] =  [0.0]
+            params['slices_other_beam_Sigma_11'] = [bb_df.loc[eename, 'other_Sigma_11']]
+            params['slices_other_beam_Sigma_12'] = [bb_df.loc[eename, 'other_Sigma_12']]
+            params['slices_other_beam_Sigma_13'] = [bb_df.loc[eename, 'other_Sigma_13']]
+            params['slices_other_beam_Sigma_14'] = [bb_df.loc[eename, 'other_Sigma_14']]
+            params['slices_other_beam_Sigma_22'] = [bb_df.loc[eename, 'other_Sigma_22']]
+            params['slices_other_beam_Sigma_23'] = [bb_df.loc[eename, 'other_Sigma_23']]
+            params['slices_other_beam_Sigma_24'] = [bb_df.loc[eename, 'other_Sigma_24']]
+            params['slices_other_beam_Sigma_33'] = [bb_df.loc[eename, 'other_Sigma_33']]
+            params['slices_other_beam_Sigma_34'] = [bb_df.loc[eename, 'other_Sigma_34']]
+            params['slices_other_beam_Sigma_44'] = [bb_df.loc[eename, 'other_Sigma_44']]
 
             if not (bb_coupling):
-                params['sigma_13'] =  0.0
-                params['sigma_14'] =  0.0
-                params['sigma_23'] =  0.0
-                params['sigma_24'] =  0.0
-
-            params["x_co"] = 0
-            params["px_co"] = 0
-            params["y_co"] = 0
-            params["py_co"] = 0
-            params["zeta_co"] = 0
-            params["delta_co"] = 0
-            params["d_x"] = 0
-            params["d_px"] = 0
-            params["d_y"] = 0
-            params["d_py"] = 0
-            params["d_zeta"] = 0
-            params["d_delta"] = 0
-
-            newee = xf.BeamBeamBiGaussian3D(old_interface=params)
+                params['slices_other_beam_Sigma_13'] = [0.0]
+                params['slices_other_beam_Sigma_14'] = [0.0]
+                params['slices_other_beam_Sigma_23'] = [0.0]
+                params['slices_other_beam_Sigma_24'] = [0.0]
+
+            newee = xf.BeamBeamBiGaussian3D(**params)
             line.element_dict[eename] = newee
 
 
diff --git a/pymask/pymasktools.py b/pymask/pymasktools.py
@@ -498,52 +498,52 @@ def save_mad_sequence_and_error(mad, seq_name, filename='lhc'):
 def _disable_beam_beam(line):
     for ee in line.elements:
         if ee.__class__.__name__.startswith('BeamBeam'):
-            ee._temp_q0 = ee.q0
-            ee.q0 = 0
+            ee._temp_q0 = ee.other_beam_q0
+            ee.other_beam_q0 = 0
             if ee.__class__.__name__ == 'BeamBeamBiGaussian2D':
-                ee._temp_d_px = ee.d_px
-                ee._temp_d_py = ee.d_py
-                ee.d_px = 0.
-                ee.d_py = 0.
+                ee._temp_d_px = ee.post_subtract_px
+                ee._temp_d_py = ee.post_subtract_py
+                ee.post_subtract_px = 0.
+                ee.post_subtract_py = 0.
             elif ee.__class__.__name__ == 'BeamBeamBiGaussian3D':
-                ee._temp_Dx_sub = ee.Dx_sub
-                ee._temp_Dpx_sub = ee.Dpx_sub
-                ee._temp_Dy_sub = ee.Dy_sub
-                ee._temp_Dpy_sub = ee.Dpy_sub
-                ee._temp_Dsigma_sub = ee.Dsigma_sub
-                ee._temp_Ddelta_sub = ee.Ddelta_sub
-                ee.Dx_sub = 0.
-                ee.Dpx_sub = 0.
-                ee.Dy_sub = 0.
-                ee.Dpy_sub = 0.
-                ee.Dsigma_sub = 0.
-                ee.Ddelta_sub = 0.
+                ee._temp_Dx_sub = ee.post_subtract_x
+                ee._temp_Dpx_sub = ee.post_subtract_px
+                ee._temp_Dy_sub = ee.post_subtract_y
+                ee._temp_Dpy_sub = ee.post_subtract_py
+                ee._temp_Dzeta_sub = ee.post_subtract_zeta
+                ee._temp_Dpzeta_sub = ee.post_subtract_pzeta
+                ee.post_subtract_x = 0.
+                ee.post_subtract_px = 0.
+                ee.post_subtract_y = 0.
+                ee.post_subtract_py = 0.
+                ee.post_subtract_zeta = 0.
+                ee.post_subtract_pzeta = 0.
             else:
                 raise ValueError('What?!')
 
 def _restore_beam_beam(line):
     for ee in line.elements:
         if ee.__class__.__name__.startswith('BeamBeam'):
-            ee.q0 = ee._temp_q0
+            ee.other_beam_q0 = ee._temp_q0
             del(ee._temp_q0)
             if ee.__class__.__name__ == 'BeamBeamBiGaussian2D':
-                ee.d_px = ee._temp_d_px
-                ee.d_py = ee._temp_d_py
+                ee.post_subtract_px = ee._temp_d_px
+                ee.post_subtract_py = ee._temp_d_py
                 del(ee._temp_d_px)
                 del(ee._temp_d_py)
             elif ee.__class__.__name__ == 'BeamBeamBiGaussian3D':
-                ee.Dx_sub = ee._temp_Dx_sub
-                ee.Dpx_sub = ee._temp_Dpx_sub
-                ee.Dy_sub = ee._temp_Dy_sub
-                ee.Dpy_sub = ee._temp_Dpy_sub
-                ee.Dsigma_sub = ee._temp_Dsigma_sub
-                ee.Ddelta_sub = ee._temp_Ddelta_sub
+                ee.post_subtract_x = ee._temp_Dx_sub
+                ee.post_subtract_px = ee._temp_Dpx_sub
+                ee.post_subtract_y = ee._temp_Dy_sub
+                ee.post_subtract_py = ee._temp_Dpy_sub
+                ee.post_subtract_zeta = ee._temp_Dzeta_sub
+                ee.post_subtract_pzeta = ee._temp_Dpzeta_sub
                 del(ee._temp_Dx_sub)
                 del(ee._temp_Dpx_sub)
                 del(ee._temp_Dy_sub)
                 del(ee._temp_Dpy_sub)
-                del(ee._temp_Dsigma_sub)
-                del(ee._temp_Ddelta_sub)
+                del(ee._temp_Dzeta_sub)
+                del(ee._temp_Dpzeta_sub)
             else:
                 raise ValueError('What?!')
 
diff --git a/python_examples/hl_lhc_collisions_python/checks_and_doc/t004_check_output_consistency.py b/python_examples/hl_lhc_collisions_python/checks_and_doc/t004_check_output_consistency.py
@@ -194,7 +194,7 @@ def prepare_line(path, input_type):
                 if kk == "length":
                     if np.abs(ee_test.hxl) + np.abs(ee_test.hyl) == 0.0:
                         continue
-                if kk == "order":
+                if kk == "order" or kk == "inv_factorial_order":
                     # Checked through bal
                     continue
                 if kk == 'knl' or kk == 'ksl' or kk == 'bal':
@@ -250,43 +250,46 @@ def prepare_line(path, input_type):
 
             # Exceptions BB4D (separations are recalculated)
             if not(strict) and isinstance(ee_test, xf.BeamBeamBiGaussian2D):
-                if kk == "x_bb":
-                    if diff_abs / dtest["sigma_x"] < 0.01: # This is neede to accommodate different leveling routines (1% difference)
+                if kk == "other_beam_shift_x":
+                    if diff_abs / np.sqrt(dtest["other_beam_Sigma_11"]) < 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:
+                if kk == "other_beam_shift_y":
+                    if diff_abs / np.sqrt(dtest["other_beam_Sigma_33"]) < 0.01:
                         continue
-                if kk == "sigma_x":
-                    if diff_rel < 1e-5:
+                if kk == "other_beam_Sigma_11":
+                    if diff_rel < (1e-5)**2:
                         continue
-                if kk == "sigma_y":
-                    if diff_rel < 1e-5:
+                if kk == "other_beam_Sigma_33":
+                    if diff_rel < (1e-5)**2:
                         continue
             if isinstance(ee_test, xf.BeamBeamBiGaussian2D):
-                if kk == 'q0' or kk == 'n_particles':
+                if kk == 'other_beam_q0' or kk == 'other_beam_num_particles':
                     # ambiguity due to old interface
-                    if np.abs(ee_test.n_particles*ee_test.q0 -
-                            ee_six.n_particles*ee_six.q0 ) < 1.: # charges
+                    if np.abs(ee_test.other_beam_num_particles*ee_test.other_beam_q0 -
+                            ee_six.other_beam_num_particles*ee_six.other_beam_q0 ) < 1.: # charges
                         continue
 
             # Exceptions BB6D (angles and separations are recalculated)
             if not(strict) and isinstance(ee_test, xf.BeamBeamBiGaussian3D):
-                if kk == "alpha":
+                if kk == "_cos_alpha":
                     if diff_abs < 10e-6:
                         continue
-                if kk == "x_co" or kk == "x_bb_co" or kk == 'delta_x':
-                    if diff_abs / np.sqrt(dtest["sigma_11"]) < 0.015:
+                if kk == "_sin_alpha":
+                    if diff_abs < 10e-6:
+                        continue
+                if kk == "ref_shift_x" or kk == "other_beam_shift_x":
+                    if diff_abs / np.sqrt(dtest["slices_other_beam_Sigma_11_star"][0]) < 0.015:
                         continue
-                if kk == "y_co" or kk == "y_bb_co" or kk == 'delta_y':
-                    if diff_abs / np.sqrt(dtest["sigma_33"]) < 0.015:
+                if kk == "ref_shift_y" or kk == "other_beam_shift_y":
+                    if diff_abs / np.sqrt(dtest["slices_other_beam_Sigma_33_star"][0]) < 0.015:
                         continue
-                if kk == "zeta_co":
+                if kk == "ref_shift_zeta":
                     if diff_abs <1e-5:
                         continue
-                if kk == "delta_co":
+                if kk == "ref_shift_pzeta":
                     if diff_abs <1e-5:
                         continue
-                if kk == "px_co" or kk == 'py_co':
+                if kk == "ref_shift_px" or kk == 'ref_shift_py':
                     if diff_abs <30e-9:
                         continue
             if isinstance(ee_test, xt.XYShift):
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
@@ -28,19 +28,19 @@
 phi_weak = Phi
 phi_c_weak = Phi_c
 
-# B4 ip5
-beam_track = 'b4'
-ip_choice = 5
-plane = 'y'
-phi_weak = Phi
-phi_c_weak = Phi_c
-
-# B4 ip1
-beam_track = 'b4'
-ip_choice = 1
-plane = 'x'
-phi_weak = -Phi
-phi_c_weak = -Phi_c
+## B4 ip5
+#beam_track = 'b4'
+#ip_choice = 5
+#plane = 'y'
+#phi_weak = Phi
+#phi_c_weak = Phi_c
+#
+## B4 ip1
+#beam_track = 'b4'
+#ip_choice = 1
+#plane = 'x'
+#phi_weak = -Phi
+#phi_c_weak = -Phi_c
 
 phi_strong = -phi_weak
 phi_c_strong = -phi_c_weak
@@ -105,7 +105,8 @@ def prepare_line(path, input_type):
 R1_orbit = phi_weak * s_rel
 
 axcrab.plot(s_rel, np.array(
-    [getattr(bb, f'delta_{plane}') for bb in bb_elems])+R1_orbit,
+    #[getattr(bb, f'delta_{plane}') for bb in bb_elems])+R1_orbit,
+    [getattr(bb, f'other_beam_shift_{plane}') for bb in bb_elems])+R1_orbit,
     'o', color = 'r', alpha=.5, label='strong xsuite')
 plt.plot(s_rel, R_no_crab + R_crab, '*', color='darkred', label='strong formula')
 
@@ -125,7 +126,8 @@ def prepare_line(path, input_type):
 bb_all, _ = ltest.get_elements_of_type([xf.BeamBeamBiGaussian2D,
                                         xf.BeamBeamBiGaussian3D])
 for bb in bb_all:
-    bb.q0 = 0
+    #bb.q0 = 0
+    bb.other_beam_q0 = 0
 
 # # Switch off all beam-beam lenses
 crabs, crab_names = ltest.get_elements_of_type([xt.RFMultipole])
