File indexing completed on 2025-01-18 09:10:44
0001 import sys
0002 import argparse
0003 import numpy as np
0004 import uproot
0005
0006
0007 parser = argparse.ArgumentParser()
0008 parser.add_argument("a")
0009 parser.add_argument("b")
0010 parser.add_argument("--event-nr", default="event_nr")
0011 parser.add_argument("--fail-fast", action="store_true")
0012 args = parser.parse_args()
0013
0014 a_data = uproot.open(args.a)
0015 b_data = uproot.open(args.b)
0016 event_nr = args.event_nr
0017 fail_fast = args.fail_fast
0018
0019 a_sort_index = np.argsort(a_data[event_nr].array(library="np"), kind="stable")
0020 b_sort_index = np.argsort(b_data[event_nr].array(library="np"), kind="stable")
0021
0022 np.set_printoptions(linewidth=np.inf)
0023
0024
0025 def cmp(a, b):
0026 if np.isscalar(a):
0027 if a != b:
0028 return False
0029 elif type(a) == np.ndarray:
0030 if not np.array_equal(a, b, equal_nan=True):
0031 return False
0032 else:
0033 for aa, bb in zip(a, b):
0034 if not cmp(aa, bb):
0035 return False
0036 return True
0037
0038
0039 failed_events = set()
0040 failed_keys = set()
0041 for key in a_data.keys():
0042 if key == event_nr:
0043 continue
0044
0045 a_vals = a_data[key].array(library="np")
0046 a_vals = a_vals[a_sort_index]
0047
0048 b_vals = b_data[key].array(library="np")
0049 b_vals = b_vals[b_sort_index]
0050
0051 for event, (a, b) in enumerate(zip(a_vals, b_vals)):
0052 if not cmp(a, b):
0053 print(f"event {event} failed for {key}")
0054 print(f"a {a}")
0055 print(f"b {b}")
0056 print()
0057
0058 if fail_fast:
0059 sys.exit(1)
0060 failed_events.add(event)
0061 failed_keys.add(key)
0062
0063 if failed_events or failed_keys:
0064 print("summary")
0065 print("failed events: " + " ".join(map(str, sorted(failed_events))))
0066 print("failed keys: " + " ".join(failed_keys))
0067 sys.exit(1)
0068
0069 sys.exit(0)
