Merge branch 'master' of github.com:lhcopt/lhcmask

pymask/pymasktools.py

@@ -408,13 +408,13 @@ def generate_xsuite_line(mad, seq_name, bb_df,
         folder_name=None, skip_mad_use=False,
         prepare_line_for_xtrack=True,
         steps_for_finite_diffs={'dx': 1e-8, 'dpx': 1e-11,
-            'dy': 1e-8, 'dpy': 1e-11, 'dzeta': 1e-7, 'ddelta': 1e-8}, 
+            'dy': 1e-8, 'dpy': 1e-11, 'dzeta': 1e-7, 'ddelta': 1e-8},
         deferred_expressions=True):
 
     # Build xsuite model
     print('Start building xtrack line...')
     line = xt.Line.from_madx_sequence(
-        mad.sequence[seq_name], apply_madx_errors=True, 
+        mad.sequence[seq_name], apply_madx_errors=True,
         deferred_expressions=deferred_expressions)
     print('Done building xtrack.')
 
@@ -478,7 +478,7 @@ def generate_xsuite_line(mad, seq_name, bb_df,
             with open(folder_name +
                     '/line_bb_for_tracking.json', 'w') as fid:
                 json.dump(line_bb_for_tracking_dict, fid, cls=JEncoder)
-
+        return tracker, line_bb_for_tracking_dict 
 
 def save_mad_sequence_and_error(mad, seq_name, filename='lhc'):
     mad.select(flag="error",clear=True)

python_examples/hl_lhc_collisions_python/001_reload.py

@@ -1,7 +1,7 @@
 import sys
 import numpy as np
 import xtrack as xt
-
+import pandas as pd
 sys.path.append('./modules')
 import pymask as pm
 
@@ -15,3 +15,23 @@
 mad.set(format=".15g")
 mad.twiss(rmatrix = True)
 
+sequence_to_track = 'lhcb1'
+bb_df_track = pd.read_pickle('bb_df_track.pickle')
+optics_and_co_at_start_ring_from_madx = pm.get_optics_and_orbit_at_start_ring(
+    mad, sequence_to_track, skip_mad_use=True)
+#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 xtrack line #                                                           
+########################                                                           
+#if enable_bb_legacy:
+#    print('xtrack line is not generated with bb legacy macros')
+#    else:
+
+tracker, line_bb_for_tracking_dict = pm.generate_xsuite_line(mad,
+                        sequence_to_track, bb_df_track,
+                        optics_and_co_at_start_ring_from_madx,
+                        folder_name = './xsuite_lines',
+                        skip_mad_use=True,
+                        prepare_line_for_xtrack=True)

python_examples/hl_lhc_collisions_python/checks_and_doc/t004_check_output_consistency.py

@@ -30,7 +30,7 @@
         'path_ref': '../',
         'type_ref': 'sixtrack',
         'rtol': 4e-7,
-        'atol': 1e-100,
+        'atol': 1e-13,
         'strict': True,
     }
 ]
@@ -153,16 +153,22 @@ def prepare_line(path, input_type):
             # 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]
+            if not np.isscalar(val_ref):
+                if len(val_ref) != len(val_test):
+                    diff_rel = 100
                 else:
-                    diff_rel = norm(np.array(val_test) - np.array(val_ref)) / norm(val_test)
-            except ZeroDivisionError:
-                diff_rel = 100.0
+                    for iiii, (vvr, vvt) in enumerate(list(zip(val_ref, val_test))):
+                        if not np.isclose(vvr, vvt, atol=atol, rtol=rtol):
+                            print(f'Issue found on `{kk}[{iiii}]`')
+                            diff_rel = 1000
+                        else:
+                            diff_rel = 0
+            else:
+                if val_ref > 0:
+                    diff_rel = np.abs((val_test - val_ref)/val_ref)
+                else:
+                    diff_rel = 100
+
             if diff_rel < rtol:
                 continue
 
@@ -209,12 +215,15 @@ def prepare_line(path, input_type):
                     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)
+                        for iiii, (vvr, vvt) in enumerate(list(zip(val_ref, val_test))):
+                            if not np.isclose(vvr, vvt, atol=atol, rtol=rtol):
+                                print(f'Issue found on `{kk}[{iiii}]`')
+                                diff_rel = 1000
+                                break
+                            else:
+                                diff_rel = 0
                         if diff_rel < rtol:
                             continue
-                        if diff_abs < atol:
-                            continue
 
             # Exception: correctors involved in lumi leveling
             passed_corr = False
@@ -229,9 +238,13 @@ def prepare_line(path, input_type):
                 '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 nn_corr in nn_test and kk in ['knl','ksl','bal']:
+                    assert len(val_ref)<=2
+                    assert len(val_test)<=2
+                    diff_rel = norm(val_test-val_ref)/norm(val_ref)
+                    passed_corr = (diff_rel < 1e-2)
+                    if passed_corr:
+                        break
             if not(strict) and  passed_corr:
                 continue
 
@@ -276,7 +289,14 @@ def prepare_line(path, input_type):
                 if kk == "px_co" or kk == 'py_co':
                     if diff_abs <30e-9:
                         continue
-
+            if isinstance(ee_test, xt.XYShift):
+                if kk in ['dx','dy']:
+                    if diff_abs <1e-9:
+                        continue
+            if isinstance(ee_test, xt.SRotation):
+                if kk in ['sin_z', 'cos_z', 'angle']:
+                    if diff_abs <1e-9:
+                        continue
             # If it got here it means that no condition above is met
             raise ValueError("Too large discrepancy!")
     print(

python_examples/hl_lhc_collisions_python/checks_and_doc/t007_check_orbit_and_lin_normal_form.py

@@ -21,7 +21,7 @@
     tracker.track(particles)
 
 for nn in 'x px y py zeta delta'.split():
-    assert np.abs(getattr(particles, nn) - getattr(partCO, nn)) < 2e-11
+    assert np.abs(getattr(particles, nn) - getattr(partCO, nn)) < 3e-11
 
 WW = np.array(line_dict['WW_finite_diffs'])
 WWinv = np.array(line_dict['WWInv_finite_diffs'])
@@ -39,13 +39,14 @@
 
 particles_matched  = xp.build_particles(particle_on_co=partCO,
                                         x_norm=x_norm, px_norm=px_norm,
-                                        R_matrix=np.array(line_dict['RR_finite_diffs']))
+                                        R_matrix=np.array(line_dict['RR_finite_diffs']),
+                                        symplectify=True)
 particles_test = particles_matched.copy()
 tracker.track(particles_test, num_turns=10)
 
 i_matched = np.argmax(particles_matched.x)
 i_test = np.argmax(particles_test.x)
-assert np.abs(particles_test.x[i_test] - particles_matched.x[i_matched]) < 1e-6
+assert np.abs(particles_test.x[i_test] - particles_matched.x[i_matched]) < 2e-6
 assert np.abs(particles_test.px[i_test] - particles_matched.px[i_matched]) < 1e-7