Some fixes in examples

Author: sterbini

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.')
 

python_examples/hl_lhc_collisions_python/checks_and_doc/t004_check_output_consistency.py

@@ -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 = (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

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