Make it more robust

Author: sterbini

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,
     }
 ]
@@ -238,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
 
@@ -285,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(