Tracking/track2particle/epic_analysis_base.py

0001 '''
0002     Utility functions for epic tracking analysis.
0003     See also epic_analysis_podio.py
0004     Shujie Li, Aug 2025
0005 '''
0006
0007 # %load_ext memory_profiler
0008 import numpy as np
0009 import pandas as pd
0010 import seaborn as sns
0011
0012 import awkward as ak
0013 import uproot as ur
0014 # import podio
0015 # from podio import root_io
0016
0017 import time
0018 from fnmatch import fnmatch
0019 import types
0020 from particle import Particle
0021
0022 from matplotlib.backends.backend_pdf import PdfPages
0023 from matplotlib.colors import LogNorm
0024 from matplotlib import pyplot as plt
0025 from matplotlib.gridspec import GridSpec
0026 import matplotlib.ticker as ticker
0027 import matplotlib.cm as cm
0028 import matplotlib as mpl
0029
0030 def configure_analysis_environment(
0031     apply_pandas=True,
0032     apply_matplotlib=True,
0033     apply_sns=True
0034 ):
0035     """Apply pandas/matplotlib defaults for interactive analysis."""
0036     if apply_pandas:
0037         pd.options.display.max_rows = 200
0038         pd.options.display.min_rows = 20
0039         pd.options.display.max_columns = 100
0040     if apply_matplotlib:
0041         plt.rcParams['figure.figsize'] = [8.0, 6.0]
0042         plt.rcParams['ytick.direction'] = 'in'
0043         plt.rcParams['xtick.direction'] = 'in'
0044         plt.rcParams['xaxis.labellocation'] = 'right'
0045         plt.rcParams['yaxis.labellocation'] = 'top'
0046         small_size = 10
0047         medium_size = 12
0048         bigger_size = 20
0049         plt.rc('font', size=small_size)
0050         plt.rc('axes', titlesize=medium_size)
0051         plt.rc('axes', labelsize=medium_size)
0052         plt.rc('xtick', labelsize=medium_size)
0053         plt.rc('ytick', labelsize=medium_size)
0054         plt.rc('legend', fontsize=small_size)
0055         plt.rc('figure', titlesize=bigger_size)
0056     if apply_sns:
0057         sns.set_theme(
0058             style='whitegrid',
0059             context='notebook',
0060             palette='bright',
0061             font_scale=1.0,
0062             rc={'figure.figsize': (6, 4)},
0063         )
0064
0065 # Constants
0066 deg2rad = np.pi/180.0
0067
0068 ## event source
0069 status_to_source = {
0070     1: "DIS", 2: "DIS",
0071     2001: "SR", 2002: "SR",
0072     3001: "Bremstrahlung", 3002: "Bremstrahlung",
0073     4001: "Coulomb", 4002: "Coulomb",
0074     5001: "Touschek", 5002: "Touschek",
0075     6001: "Proton beam gas", 6002: "Proton beam gas"
0076 }
0077
0078 ## track quality cuts
0079 TRACK_HIT_COUNT_MIN_MIN   = 4 ## absolute min to form a track with CKF
0080 TRACK_HIT_COUNT_MIN       = 4
0081 TRACK_MOM_MIN             = 0.2
0082 TRACK_PT_MIN              = 0.2
0083 TRACK_HIT_FRACTION_MIN    = 0.5
0084 TRACK_HIT_COUNT_GHOST_MAX = 2
0085 VERTEX_CUT_R_MAX          = 1
0086 VERTEX_CUT_Z_MAX          = 100#mm
0087
0088 # Detector geometry definitions (unchanged)
0089 barrel_range = [(30,42),(46,60),(115,130),(250,290),(400,450),(540,600),(620,655),(700,760)]
0090 barrel_name = ["L0","L1","L2","L3","L4","inner MPGD","TOF","outer MPGD"]
0091 name_sim_barrel = ["VertexBarrelHits","VertexBarrelHits","VertexBarrelHits","SiBarrelHits","SiBarrelHits","MPGDBarrelHits","TOFBarrelHits","OuterMPGDBarrelHits"]
0092 name_rec_barrel = ["SiBarrelVertexRecHits","SiBarrelVertexRecHits","SiBarrelVertexRecHits","SiBarrelTrackerRecHits","SiBarrelTrackerRecHits","MPGDBarrelRecHits","TOFBarrelRecHits","OuterMPGDBarrelRecHits"]
0093
0094 disk_range = [(-1210.0, -1190.0), (-1110.0, -1090.0),(-1055.0, -1000.0), (-860.0, -840.0),
0095  (-660.0, -640.0), (-460.0, -440.0), (-260.0, -240.0), (240.0, 260.0),
0096  (440.0, 460.0), (690.0, 710.0), (940.0, 960.0), (1150.0, 1250.0),
0097  (1480.0, 1500.0), (1600.0, 1620.0), (1840.0, 1860.0), (1865.0, 1885.0)]
0098 disk_name = ["E-MPGD Disk2","E-MPGD Disk 1","E-Si Disk 4","E-Si Disk 3","E-Si Disk 2","E-Si Disk 1","E-Si Disk 0",
0099                 "H-Si Disk 0","H-Si Disk 1","H-Si Disk 2","H-Si Disk 3","H-Si Disk 4","H-MPGD Disk 1","H-MPGD Disk 2", "H-TOF Disk1","H-TOF Disk2"]
0100 name_rec_disk = ["BackwardMPGDEndcapRecHits","BackwardMPGDEndcapRecHits",
0101                "SiEndcapTrackerRecHits","SiEndcapTrackerRecHits","SiEndcapTrackerRecHits","SiEndcapTrackerRecHits","SiEndcapTrackerRecHits","SiEndcapTrackerRecHits","SiEndcapTrackerRecHits","SiEndcapTrackerRecHits","SiEndcapTrackerRecHits","SiEndcapTrackerRecHits",
0102                "ForwardMPGDEndcapRecHits","ForwardMPGDEndcapRecHits",
0103                "TOFEndcapRecHits","TOFEndcapRecHits"]
0104 name_sim_disk = ["BackwardMPGDEndcapHits","BackwardMPGDEndcapHits",
0105                "TrackerEndcapHits","TrackerEndcapHits","TrackerEndcapHits","TrackerEndcapHits","TrackerEndcapHits","TrackerEndcapHits","TrackerEndcapHits","TrackerEndcapHits","TrackerEndcapHits","TrackerEndcapHits","ForwardMPGDEndcapHits","ForwardMPGDEndcapHits",
0106                "TOFEndcapHits","TOFEndcapHits"]
0107
0108 # ACTS geometry ID masks
0109 geo_mask_dict = {
0110     "approach": 0x0000000ff0000000,
0111     "boundary": 0x00ff000000000000,
0112     "extra": 0x00000000000000ff,
0113     "layer": 0x0000fff000000000,
0114     "sensitive": 0x000000000fffff00,
0115     "volume": 0xff00000000000000
0116 }
0117 geo_mask_values = types.MappingProxyType(geo_mask_dict)
0118
0119 # Global variables for caching
0120 COL_TABLE = {}
0121 CACHED_DATA = {}
0122
0123 def ak_flat(ak_array):
0124     return ak.to_numpy(ak.flatten(ak_array,axis=0))
0125
0126 def ak_df(ak_array):
0127     return ak.to_dataframe(ak_array)
0128
0129 def ak_hist(ak_array, **kwargs):
0130     return plt.hist(ak_flat(ak_array), **kwargs)
0131
0132 def ak_filter(br, cond, field=None):
0133     filtered = br[cond]
0134     return filtered[field] if field else filtered
0135
0136
0137 def ak_sns(ak_array, **kwargs):
0138     """Histogram helper for awkward arrays using seaborn."""
0139     if isinstance(ak_array, (tuple, list)) and len(ak_array) == 2:
0140         x_data = ak_flat(ak_array[0])
0141         y_data = ak_flat(ak_array[1])
0142         kwargs.pop('element', None)
0143         kwargs.pop('fill', None)
0144         return sns.histplot(x=x_data, y=y_data, **kwargs)
0145     return sns.histplot(ak_flat(ak_array), element="step", fill=False, **kwargs)
0146
0147 def get_pdg_info(PDG):
0148     """Get particle info from PDG code"""
0149     try:
0150         return Particle.from_pdgid(PDG)
0151     except Exception:
0152         if PDG == 9902210:
0153             return Particle.from_pdgid(2212)
0154         print(f"ERROR (get_pdg_info): unknown PDG ID {PDG}")
0155         return Particle.empty()
0156
0157 def get_geoID(geoID, name="layer"):
0158     """Extract geometry ID components"""
0159     kMask = geo_mask_values[name]
0160     shift = 0
0161     mask_temp = kMask
0162     while (mask_temp & 1) == 0:
0163         mask_temp >>= 1
0164         shift += 1
0165     return (geoID & kMask) >> shift
0166
0167
0168 def theta2eta(xx, inverse=0):
0169     """Convert theta to eta or vice versa"""
0170     if type(xx)==list:
0171         xx = np.array(xx)
0172     if inverse==1:
0173         return np.arctan((np.e)**(-xx))*2
0174     else:
0175         return -np.log(np.tan(xx/2.))
0176
0177 def select_string(strings, patterns):
0178     """Select strings matching patterns with wildcards"""
0179     if not isinstance(patterns, list):
0180         raise ValueError("The 'patterns' argument must be a list.")
0181
0182     patterns = [pattern.lower() for pattern in patterns]
0183     return [s for s in strings if any(fnmatch(s.lower(), pattern) for pattern in patterns)]
0184
0185 def read_ur(fname, tname, s3_dir="", entry_start=0, entry_stop=None, return_range=False):
0186     """Read ROOT file with uproot
0187     fname: path to file
0188     tname: tree name
0189     s3_dir: if provided, try the BNL and JLab simulation campaign locations"""
0190     if len(s3_dir) > 0:
0191         # The same campaign path is stored below different roots at BNL and JLab.
0192         servers = (
0193             'root://epicxrd1.sdcc.bnl.gov:1095//eic/',
0194             'root://dtn-eic.jlab.org//volatile/eic/',
0195             "root://epicxrd1.sdcc.bnl.gov:1095/",
0196         )
0197         remote_path = s3_dir.rstrip('/') + '/' + fname.lstrip('/')
0198         open_errors = []
0199         last_error = None
0200         for server in servers:
0201             candidate = server + remote_path
0202             try:
0203                 root_file = ur.open(candidate, timeout=5)
0204                 fname = candidate
0205                 break
0206             except OSError as error:
0207                 last_error = error
0208                 open_errors.append(f"{candidate}: {error}")
0209         else:
0210             attempted = '\n  '.join(open_errors)
0211             raise FileNotFoundError(
0212                 f"read_ur: remote file was not found on either server:\n  {attempted}"
0213             ) from last_error
0214     else:
0215         root_file = ur.open(fname)
0216
0217     tree = root_file[tname]
0218     if entry_stop is None or entry_stop == -1:
0219         entry_stop = tree.num_entries
0220     entry_stop = min(entry_stop, tree.num_entries)
0221     if entry_start < 0 or entry_start >= entry_stop:
0222         raise ValueError(f"read_ur: invalid entry range {entry_start}:{entry_stop}")
0223     print(
0224         f"read_ur: read {fname}:{tname}. "
0225         f"{tree.num_entries} events total; using [{entry_start}, {entry_stop})"
0226     )
0227     tree._entry_start = entry_start
0228     tree._entry_stop = entry_stop
0229     if return_range:
0230         return tree, entry_start, entry_stop
0231     return tree
0232
0233 def get_col_table(fname, s3_dir="", verb=0):
0234     """Get collection table from metadata"""
0235     global COL_TABLE
0236     meta = read_ur(fname, "podio_metadata", s3_dir)
0237     if "events___idTable" in meta.keys(): ## < eic-shell 25.09
0238         col_name = np.array(meta["m_names"].array()[0])
0239         col_id = np.array(meta["m_collectionIDs"].array()[0])
0240     else:
0241         col_id = get_branch_df(meta,"events___CollectionTypeInfo/events___CollectionTypeInfo.collectionID")["values"].tolist()
0242         col_name = get_branch_df(meta,"events___CollectionTypeInfo/events___CollectionTypeInfo.name")["values"].tolist()
0243
0244     COL_TABLE = {}
0245     for ii, nn in zip(col_id, col_name):
0246         if verb:
0247             print(ii, nn)
0248         COL_TABLE[ii] = nn
0249     return COL_TABLE
0250
0251
0252 # ============= BRANCH READING with ak or df =============
0253
0254 def get_branch_ak(tree, bname="", entry_start=0, entry_stop=-1,
0255                            fields_subset=None, chunk_size=1000, verb=0):
0256     """Optimized branch reading with awkward arrays"""
0257     if bname not in tree.keys():
0258         raise KeyError(f"get_branch_ak: can't find branch {bname}")
0259     if verb:
0260         print(f"Reading branch: {bname}")
0261     start_time = time.time()
0262
0263     # Determine actual entry range
0264     if entry_start == 0 and hasattr(tree, "_entry_start"):
0265         entry_start = tree._entry_start
0266     if entry_stop == -1:
0267         entry_stop = tree._entry_stop if hasattr(tree, "_entry_stop") else tree.num_entries
0268
0269     total_entries = entry_stop - entry_start
0270     if verb:
0271         print(f"Reading {total_entries} entries")
0272
0273     # For large datasets, read in chunks
0274     if total_entries > chunk_size:
0275         if verb:
0276             print(f"Using chunked reading with chunk_size={chunk_size}")
0277         all_data = []
0278
0279         for chunk_start in range(entry_start, entry_stop, chunk_size):
0280             chunk_end = min(chunk_start + chunk_size, entry_stop)
0281             if verb:
0282                 print(f"  Reading chunk: {chunk_start} to {chunk_end}")
0283
0284             chunk_data = tree[bname].array(
0285                 library="ak",
0286                 entry_start=chunk_start,
0287                 entry_stop=chunk_end
0288             )
0289             all_data.append(chunk_data)
0290
0291         # Concatenate all chunks
0292         ak_data = ak.concatenate(all_data)
0293     else:
0294         # Read all at once for smaller datasets
0295         ak_data = tree[bname].array(
0296             library="ak",
0297             entry_start=entry_start,
0298             entry_stop=entry_stop
0299         )
0300
0301     read_time = time.time()
0302     if verb:
0303         try:
0304             size_bytes = ak.nbytes(ak_data)
0305         except Exception:
0306             size_bytes = None
0307         if size_bytes is not None:
0308             print(f"Awkward read: {read_time - start_time:.2f}s ({size_bytes/1e6:.1f} MB)")
0309         else:
0310             print(f"Awkward read: {read_time - start_time:.2f}s")
0311
0312     # Rename fields by dropping branch prefix
0313     if hasattr(ak_data, 'fields'):
0314         if not ak_data.fields:
0315             ## return a single array
0316             return ak_data
0317         renamed_fields = {}
0318         for field in ak_data.fields:
0319             if "[" in field: ## drop nested array such as CentralCKFTrackParameters.covariance.covariance[21]
0320                 continue
0321             if field.startswith(f'{bname}.'):
0322                 new_name = field.replace(f'{bname}.', '')
0323                 renamed_fields[new_name] = ak_data[field]
0324             else:
0325                 renamed_fields[field] = ak_data[field]
0326
0327         ak_data = ak.zip(renamed_fields)
0328         if verb:
0329             print(f"Renamed {len(renamed_fields)} fields")
0330
0331     # Extract subset of fields if specified
0332     if fields_subset and hasattr(ak_data, 'fields'):
0333         subset_data = {}
0334         for field in fields_subset:
0335             if field in ak_data.fields:
0336                 subset_data[field] = ak_data[field]
0337         ak_data = ak.zip(subset_data)
0338         if verb:
0339             print(f"Extracted subset: {fields_subset}")
0340
0341     total_time = time.time()
0342     if verb:
0343             print(f"Total time: {total_time - start_time:.2f}s")
0344
0345     return ak_data
0346
0347
0348 def get_branch_df(tree, bname="", entry_start=0, entry_stop=-1,
0349                            fields_subset=None, chunk_size=1000, verb=0):
0350     """Get branch as DataFrame when needed (for compatibility)"""
0351     ak_data = get_branch_ak(tree, bname, entry_start, entry_stop,
0352                                      fields_subset, chunk_size)
0353
0354     if verb:
0355         print("Converting to DataFrame...")
0356     start_time = time.time()
0357
0358     try:
0359         df = ak.to_dataframe(ak_data)
0360         df = df.reset_index()  # Flatten the multi-index
0361
0362         convert_time = time.time() - start_time
0363         if verb:
0364             print(f"DataFrame conversion: {convert_time:.2f}s")
0365             print(f"DataFrame shape: {df.shape}")
0366
0367         return df
0368
0369     except Exception as e:
0370         print(f"DataFrame conversion failed: {e}")
0371         print("Returning awkward array instead")
0372         return ak_data
0373
0374 def get_part(tree, entry_start=0, entry_stop=-1, chunk_size=1000, kprimary=1):
0375     """MC particles reading, return ak with calculated eta, momentum etc"""
0376     # print("Reading MC particles as akward arrays...")
0377
0378     # Read as awkward array
0379     ak_data = get_branch_ak(tree, "MCParticles", entry_start, entry_stop,
0380                                      chunk_size=chunk_size)
0381     # Preserve original per-event particle index so downstream code can
0382     # match podio particle.id().index even after filtering.
0383     orig_subentry = ak.local_index(ak_data.PDG, axis=1)
0384
0385     if kprimary:
0386         print("Select all primary particles with generatorStatus==x001 or x002")
0387         primary_mask = (ak_data.generatorStatus % 1000 == 1) | (ak_data.generatorStatus % 1000 == 2)
0388         ak_data = ak_data[primary_mask]
0389         orig_subentry = orig_subentry[primary_mask]
0390     # Compute momentum quantities vectorized
0391     px = ak_data["momentum.x"]
0392     py = ak_data["momentum.y"]
0393     pz = ak_data["momentum.z"]
0394
0395     p_mag = np.sqrt(px**2 + py**2 + pz**2)
0396     safe_p_mag = ak.where(p_mag != 0, p_mag, np.nan)
0397     theta = np.arccos(pz / safe_p_mag)
0398     phi = np.arctan2(py, px)
0399     eta = -np.log(np.tan(theta / 2.0))
0400     pt = p_mag * np.sin(theta)
0401
0402     # Compute vertex quantities
0403     vx = ak_data["vertex.x"]
0404     vy = ak_data["vertex.y"]
0405     vz = ak_data["vertex.z"]
0406     vertex_r = np.sqrt(vx**2 + vy**2)
0407     vertex_dist = np.sqrt(vertex_r**2 + vz**2)
0408
0409     # Compute endpoint quantities
0410     ex = ak_data["endpoint.x"]
0411     ey = ak_data["endpoint.y"]
0412     endpoint_r = np.sqrt(ex**2 + ey**2)
0413
0414     # Get PDG names (this is slower, so we do it efficiently)
0415     pdg_codes = ak.to_numpy(ak.flatten(ak_data.PDG))
0416     unique_pdgs = np.unique(pdg_codes)
0417
0418     # Create PDG name mapping for unique values only
0419     pdg_name_map = {}
0420     for pdg in unique_pdgs:
0421         pdg_name_map[pdg] = get_pdg_info(pdg).name
0422
0423     # Apply mapping vectorized
0424     pdg_names = ak.unflatten(
0425         np.array([pdg_name_map[pdg] for pdg in pdg_codes]),
0426         ak.num(ak_data.PDG)
0427     )
0428
0429     # Create new awkward array with selected quantities
0430     my_field=['PDG', 'generatorStatus',  'charge', 'time', 'mass',
0431        'vertex.x', 'vertex.y', 'vertex.z', 'endpoint.x', 'endpoint.y',
0432        'endpoint.z', 'momentum.x', 'momentum.y', 'momentum.z']
0433     enhanced_data = ak.zip({
0434         # Original fields
0435         **{field: ak_data[field] for field in my_field},
0436         # Derived quantities
0437         'mom': p_mag,
0438         'theta': theta,
0439         'phi': phi,
0440         'eta': eta,
0441         'pt': pt,
0442         'vertex_r': vertex_r,
0443         'vertex_dist': vertex_dist,
0444         'endpoint_r': endpoint_r,
0445         'pdg_name': pdg_names,
0446         'orig_subentry': orig_subentry
0447     })
0448
0449     return enhanced_data
0450
0451 def get_params(tree, bname="CentralCKFTrackParameters", entry_start=0, entry_stop=-1, chunk_size=1000):
0452     """Track Parameters reading, return ak with calculated eta, mom, pt"""
0453
0454     # Read as awkward array
0455     ak_data = get_branch_ak(tree, bname, entry_start, entry_stop,
0456                                      chunk_size=chunk_size)
0457     eta = theta2eta(ak_data.theta)
0458     mom = abs(1.0 / ak_data.qOverP)
0459     pt = abs(mom * np.sin(ak_data.theta))
0460     ak_data = ak.with_field(ak_data, eta, "eta")
0461     ak_data = ak.with_field(ak_data, mom, "mom")
0462     ak_data = ak.with_field(ak_data, pt, "pt")
0463     return ak_data
0464
0465 def get_branches(trees,bname):
0466     df=pd.DataFrame()
0467     for tree in trees:
0468         if bname=="MCParticles":
0469             dff = get_part(tree)
0470         else:
0471             dff=get_branch_df(tree,bname)
0472         df = pd.concat([df,dff],ignore_index=True)
0473
0474         ## add a new counter (event_id)
0475         event_id = []
0476         current_id = -1
0477         prev_entry = None
0478         for e in df["entry"]:
0479             if prev_entry is None:
0480                 current_id += 1
0481             elif e != prev_entry:
0482                 current_id += 1
0483             event_id.append(current_id)
0484             prev_entry = e
0485         df["event_id"] = event_id
0486
0487     return df
0488
0489
0490 def get_collections(tree, bname='', kflatten=1):
0491     """Extract collections that a given branch pointed to (one to one/many relation)"""
0492
0493     if not COL_TABLE:
0494         raise RuntimeError("COL_TABLE not populated. Call get_col_table() first.")
0495     # Use the optimized awkward array reader
0496     br = get_branch_ak(tree, bname) if kflatten else get_branch_df(tree, bname, chunk_size=1000)
0497
0498     # Convert to DataFrame to check for collectionID
0499     if hasattr(br, 'fields') and 'collectionID' in br.fields:
0500         # Get unique collection IDs efficiently
0501         colID = np.unique(ak.to_numpy(ak.flatten(br.collectionID)))
0502         collections = {}
0503
0504         print(f"Loading {len(colID)} collections...")
0505         for ii in colID:
0506             if ii in COL_TABLE:
0507                 # Use optimized reader for each collection
0508                 collections[ii] = get_branch_ak(tree, COL_TABLE[ii]) if kflatten else get_branch_df(tree, COL_TABLE[ii], chunk_size=1000)
0509             else:
0510                 print(f"Warning: Collection ID {ii} not found in COL_TABLE")
0511
0512         return collections
0513     else:
0514         # Convert to DataFrame to check columns if it's awkward array
0515         if hasattr(br, 'fields'):
0516             br_df = ak.to_dataframe(br).reset_index()
0517         else:
0518             br_df = br
0519
0520         if "collectionID" in br_df.columns:
0521             colID = br_df.collectionID.unique()
0522             collections = {}
0523             for ii in colID:
0524                 if ii in COL_TABLE:
0525                     collections[ii] = get_branch_df(tree, COL_TABLE[ii], chunk_size=1000)
0526             return collections
0527         else:
0528             print("ERROR(get_collections):", bname, "is not a relation.")
0529             return 0
0530
0531 def get_relation(tree, b_name, v_name):
0532     """Get relation or vector members with index"""
0533     print(f"Processing relation: {b_name}.{v_name}")
0534
0535     # Read main branch as DataFrame for index operations
0536     br = get_branch_df(tree, b_name, chunk_size=1000)
0537
0538     # Check if the required columns exist
0539     begin_col = v_name + "_begin"
0540     end_col = v_name + "_end"
0541
0542     if begin_col not in br.columns or end_col not in br.columns:
0543         print(f"ERROR(get_relation): {begin_col} or {end_col} not found in {b_name}")
0544         return 0
0545
0546     loc1 = br.columns.get_loc(begin_col)
0547     loc2 = br.columns.get_loc(end_col)
0548     in_name = "_" + b_name + "_" + v_name
0549
0550     # Read the relation/vector data
0551     try:
0552         app = get_branch_df(tree, in_name, chunk_size=1000)
0553     except:
0554         print(f"ERROR(get_relation): Cannot read branch {in_name}")
0555         return 0
0556
0557     if not isinstance(app, pd.DataFrame):
0558         print(f"ERROR(get_relation): {in_name} is not a valid DataFrame")
0559         return 0
0560
0561     # Vector of float values
0562     if len(app.columns) == 3 and app.columns[2] == "values":
0563         print("Processing vector values...")
0564         l_val = []
0565         l_ind = []
0566         # Vectorized approach for better performance
0567         for row in br.itertuples(index=False):
0568             l_ind.append(row[0])
0569
0570             i1, i2 = row[loc1], row[loc2]
0571             if i1 == i2:  # empty
0572                 l_val.append(np.array([]))
0573             else:
0574                 # Use entry to locate the correct event
0575                 event_data = app[app['entry'] == row.entry]
0576                 if not event_data.empty:
0577                     v_row = np.array(event_data['values'])[i1:i2]
0578                     l_val.append(v_row)
0579
0580                 else:
0581                     l_val.append(np.array([]))
0582
0583         return pd.DataFrame({"entry": l_ind, "values": l_val})
0584
0585     # Relations with collectionID and index
0586     elif len(app.columns) > 1 and 'index' in app.columns and 'collectionID' in app.columns:
0587         print("Processing relations...")
0588         l_index = []
0589         l_id = []
0590
0591         # Group app by entry for faster lookup
0592         app_grouped = app.groupby('entry')
0593
0594         for row in br.itertuples(index=False):
0595             i1, i2 = row[loc1], row[loc2]
0596             if i1 == i2:  # empty
0597                 l_index.append(np.array([]))
0598                 l_id.append(np.array([]))
0599             else:
0600                 # Get the event data
0601                 if row.entry in app_grouped.groups:
0602                     v_row = app_grouped.get_group(row.entry)
0603                     if len(v_row) >= i2:  # Check bounds
0604                         l1 = np.array(v_row["index"])[i1:i2]
0605                         l2 = np.array(v_row["collectionID"])[i1:i2]
0606                         # l2 = np.array(v_row["collectionID"].iloc[0])[i1:i2] if hasattr(v_row["collectionID"].iloc[0], '__iter__') else np.array(v_row["collectionID"])[i1:i2]
0607
0608                         l_index.append(l1)
0609                         l_id.append(l2)
0610                     else:
0611                         l_index.append(np.array([]))
0612                         l_id.append(np.array([]))
0613                 else:
0614                     l_index.append(np.array([]))
0615                     l_id.append(np.array([]))
0616
0617         return pd.DataFrame({"index": l_index, "collectionID": l_id})
0618
0619     else:
0620         print("ERROR(get_relation): Invalid vector or relation member structure")
0621         print(f"Columns found: {app.columns.tolist()}")
0622         return 0
0623
0624 def get_branch_relation(tree, branch_name="CentralCKFTrajectories", relation_name="measurements_deprecated", relation_variables=["*"]):
0625     """Get relation or vector members appended to the original branch with optimization"""
0626     print(f"Processing branch relation: {branch_name}.{relation_name}")
0627     if not COL_TABLE:
0628         raise RuntimeError("COL_TABLE not populated. Call get_col_table() first.")
0629     # Read main branch
0630     br = get_branch_df(tree, branch_name, chunk_size=1000)
0631     df = get_relation(tree, branch_name, relation_name)
0632
0633     if not isinstance(df, pd.DataFrame):
0634         print("ERROR (get_branch_relation): please provide a valid relation name.")
0635         return br, None
0636
0637     # Handle vector values case
0638     if len(df.columns) == 1 and df.columns[0] == "values":
0639         return br, df["values"]
0640
0641     # Handle relations case
0642     br[relation_name + "_index"] = df["index"]
0643     br[relation_name + "_colID"] = df["collectionID"]
0644
0645     # Return early if no relation variables requested
0646     if len(relation_variables) == 0:
0647         br = br.explode([relation_name + "_index", relation_name + "_colID"]).reset_index(drop=True)
0648         return br, None
0649
0650     # Prepare collections
0651     in_name = "_" + branch_name + "_" + relation_name
0652     print("Loading collections...")
0653     collections = get_collections(tree, in_name, 0)
0654
0655     if not collections:
0656         print(f"ERROR: No collections {in_name} found")
0657         return br, None
0658
0659     # Process relations efficiently
0660     print("Processing relation data...")
0661     l_relations = []
0662     l_collection_names = []
0663     loc1 = br.columns.get_loc(relation_name + "_index")
0664     loc2 = br.columns.get_loc(relation_name + "_colID")
0665
0666     # Pre-convert collections to grouped DataFrames for faster access
0667     collections_grouped = {}
0668     sample_columns = None
0669
0670     for col_id, col_data in collections.items():
0671         if hasattr(col_data, 'fields'):
0672             col_df = ak.to_dataframe(col_data).reset_index()
0673         else:
0674             col_df = col_data
0675         collections_grouped[col_id] = col_df.groupby('entry')
0676         if sample_columns is None:
0677             sample_columns = col_df.columns
0678
0679     # Process each row
0680     for row in br.itertuples(index=False):
0681         ind = row[loc1]
0682         col = row[loc2]
0683         if len(ind) == 0:  # empty relation
0684             l_collection_names.append(None)
0685             # Add NaN row with correct number of columns
0686             if sample_columns is not None:
0687                 l_relations.append([np.nan] * (len(sample_columns) - 2))  # -2 for entry, subentry
0688             else:
0689                 l_relations.append([np.nan])
0690         else:
0691             for ii, cc in zip(ind, col):
0692                 if cc in COL_TABLE:
0693                     l_collection_names.append(COL_TABLE[cc])
0694
0695                     # Get data from pre-grouped collections
0696                     if cc in collections_grouped and row.entry in collections_grouped[cc].groups:
0697                         event_data = collections_grouped[cc].get_group(row.entry)
0698                         if ii < len(event_data):
0699                             # Get the row data, excluding entry and subentry columns
0700                             row_data = event_data.iloc[ii]
0701                             filtered_data = [row_data[col] for col in row_data.index if col not in ['entry', 'subentry']]
0702                             l_relations.append(filtered_data)
0703                         else:
0704                             l_relations.append([np.nan] * (len(sample_columns) - 2))
0705                     else:
0706                         l_relations.append([np.nan] * (len(sample_columns) - 2))
0707                 else:
0708                     l_collection_names.append(f"Unknown_{cc}")
0709                     l_relations.append([np.nan] * (len(sample_columns) - 2))
0710
0711     # Explode the main branch
0712     br = br.explode([relation_name + "_index", relation_name + "_colID"]).reset_index(drop=True)
0713
0714     # Create the additional DataFrame
0715     if sample_columns is not None:
0716         column_names = [col for col in sample_columns if col not in ['entry', 'subentry']]
0717     else:
0718         column_names = ['value']  # fallback
0719
0720     df_add = pd.DataFrame(l_relations, columns=column_names)
0721
0722     # Filter columns based on relation_variables
0723     if relation_variables != ["*"]:
0724         available_columns = select_string(column_names, relation_variables)
0725         df_add = df_add[available_columns]
0726
0727     # Add metadata
0728     df_add[relation_name + "_colName"] = l_collection_names
0729
0730     # Add entry and subentry columns at the beginning
0731     df_add.insert(0, 'entry', br['entry'])
0732     df_add.insert(1, 'subentry', br['subentry'])
0733
0734     print(f"Completed processing. Result shape: {df_add.shape}")
0735     return br, df_add
0736
0737 def get_traj_relations(tree,bname="CentralCKFTrajectories",l_var=["measurementChi2", "outlierChi2", "trackParameters", "measurements_deprecated", "outliers_deprecated"]):
0738     # l_var   = ["measurementChi2", "outlierChi2", "trackParameters", "measurements_deprecated", "outliers_deprecated"]#,"seed"]
0739     br      = get_branch_df(tree,bname)
0740     print("get_traj_relations: accessing the following vector members:")
0741     for vv in l_var:
0742         print(vv)
0743         a     = get_relation(tree,bname,vv)
0744         if not isinstance(a, pd.DataFrame):
0745             print(f"WARNING(get_traj_relations): {bname}.{vv} returned no data")
0746             continue
0747         for cc in a.columns:
0748             if cc=='values':
0749                 br[vv]=a[cc]
0750                 break
0751             elif cc=='index':
0752                 br[vv+'_index']=a[cc]
0753             elif cc=='collectionID':
0754                 br[vv+'_colID']=a[cc]
0755             else:
0756                 print('WARNING: invalid column ',cc,' in ',bname,'_',vv)
0757     return br
0758
0759
0760 def get_traj_hits_particle(tree, traj_name = 'CentralCKFTrajectories', measurement_name='measurements_deprecated'):
0761     """Trace trajectory measurements to reconstructed hits and MC particles.
0762
0763     ``CentralTrackerMeasurements`` is a PODIO subset collection in current
0764     EICrecon output. Its ``*_objIdx`` branch points to the physical measurement
0765     collections, such as ``CentralWithoutTOFTrackerMeasurements`` and the TOF
0766     cluster measurements. Resolve those references before following each
0767     measurement's ``hits`` relation.
0768
0769     A particle index of -1 means that the hit has no valid MC-particle
0770     reference, as expected for noise hits.
0771     """
0772     if not COL_TABLE:
0773         raise RuntimeError("COL_TABLE not populated. Call get_col_table() first.")
0774
0775     tree_branches = set(tree.keys())
0776     measurement_collection = "CentralTrackerMeasurements"
0777     subset_branch = f"{measurement_collection}_objIdx"
0778     measurement_hit_frames = []
0779
0780     if subset_branch in tree_branches:
0781         # Each subset entry stores the collection ID and object index of the
0782         # physical Measurement2D object. Discover the source collections from
0783         # those references instead of looking for a data branch named
0784         # CentralTrackerMeasurements.
0785         subset_references = get_branch_df(tree, subset_branch)
0786         source_collection_ids = pd.unique(subset_references["collectionID"])
0787
0788         for raw_collection_id in source_collection_ids:
0789             if pd.isna(raw_collection_id):
0790                 continue
0791             collection_id = int(raw_collection_id)
0792             if collection_id in (0, 0xFFFFFFFF):
0793                 continue
0794
0795             source_name = COL_TABLE.get(collection_id)
0796             if source_name is None:
0797                 print(
0798                     f"WARNING: measurement collection ID {collection_id} "
0799                     "is not in COL_TABLE"
0800                 )
0801                 continue
0802             if source_name not in tree_branches or f"_{source_name}_hits" not in tree_branches:
0803                 print(f"WARNING: {source_name} does not contain a readable hits relation")
0804                 continue
0805
0806             _, source_hits = get_branch_relation(
0807                 tree,
0808                 branch_name=source_name,
0809                 relation_name="hits",
0810             )
0811             if source_hits is None:
0812                 print(
0813                     f"WARNING: could not resolve the hits referenced by {source_name}. "
0814                     "The referenced reconstructed-hit collection may be missing "
0815                     "from the input file."
0816                 )
0817                 continue
0818
0819             source_hits = source_hits.copy()
0820             source_hits["measurement_colID"] = collection_id
0821             source_hits["measurement_colName"] = source_name
0822             measurement_hit_frames.append(source_hits)
0823     else:
0824         # Backward compatibility for older files where
0825         # CentralTrackerMeasurements was stored as a normal collection.
0826         _, source_hits = get_branch_relation(
0827             tree,
0828             branch_name=measurement_collection,
0829             relation_name="hits",
0830         )
0831         if source_hits is not None:
0832             measurement_ids = [
0833                 collection_id
0834                 for collection_id, name in COL_TABLE.items()
0835                 if name == measurement_collection
0836             ]
0837             source_hits = source_hits.copy()
0838             source_hits["measurement_colID"] = (
0839                 int(measurement_ids[0]) if measurement_ids else -1
0840             )
0841             source_hits["measurement_colName"] = measurement_collection
0842             measurement_hit_frames.append(source_hits)
0843
0844     if measurement_hit_frames:
0845         df_hits = pd.concat(measurement_hit_frames, ignore_index=True)
0846     else:
0847         df_hits = pd.DataFrame(
0848             columns=[
0849                 "entry",
0850                 "subentry",
0851                 "hits_colName",
0852                 "cellID",
0853                 "measurement_colID",
0854                 "measurement_colName",
0855             ]
0856         )
0857
0858     # Match reconstructed tracker hits to simulated hits by event and cell ID.
0859     # This is exact for the current one-to-one tracker digitization. The shared
0860     # TOF reconstructed-hit names are aliases for the corresponding TOF sim
0861     # hit collections.
0862     name_sim_tracker = name_sim_barrel + name_sim_disk + ["B0TrackerHits"]
0863     name_rec_tracker = name_rec_barrel + name_rec_disk + ["B0TrackerRecHits"]
0864     rec_to_sim_hit = dict(zip(name_rec_tracker, name_sim_tracker))
0865     rec_to_sim_hit.update({
0866         "TOFBarrelSharedRecHits": "TOFBarrelHits",
0867         "TOFEndcapSharedRecHits": "TOFEndcapHits",
0868     })
0869
0870     df_hits["particle_index"] = -1
0871     sim_particle_lookups = {}
0872
0873     if not df_hits.empty and "cellID" in df_hits.columns:
0874         for rec_hit_name in pd.unique(df_hits["hits_colName"].dropna()):
0875             sim_hit_name = rec_to_sim_hit.get(rec_hit_name)
0876             if sim_hit_name is None:
0877                 print(f"WARNING: {rec_hit_name} is not a recognized tracker hit collection")
0878                 continue
0879
0880             if sim_hit_name not in sim_particle_lookups:
0881                 particle_relation_name = f"_{sim_hit_name}_particle"
0882                 if (
0883                     sim_hit_name not in tree_branches
0884                     or particle_relation_name not in tree_branches
0885                 ):
0886                     print(f"WARNING: {particle_relation_name} is not a branch")
0887                     sim_particle_lookups[sim_hit_name] = pd.DataFrame()
0888                     continue
0889
0890                 sim_hits = get_branch_df(tree, sim_hit_name)
0891                 particle_relations = get_branch_df(tree, particle_relation_name)
0892
0893                 particle_indices = pd.to_numeric(
0894                     particle_relations["index"], errors="coerce"
0895                 ).fillna(-1).astype(np.int64)
0896                 particle_collection_names = particle_relations["collectionID"].map(
0897                     lambda collection_id: COL_TABLE.get(int(collection_id))
0898                     if pd.notna(collection_id)
0899                     else None
0900                 )
0901                 valid_particle_reference = (
0902                     particle_collection_names.eq("MCParticles")
0903                     & particle_indices.ge(0)
0904                     & particle_indices.ne(0xFFFFFFFF)
0905                 )
0906
0907                 particle_relations = particle_relations[
0908                     ["entry", "subentry"]
0909                 ].copy()
0910                 particle_relations["particle_index"] = np.where(
0911                     valid_particle_reference,
0912                     particle_indices,
0913                     -1,
0914                 )
0915
0916                 sim_hit_lookup = sim_hits[
0917                     ["entry", "subentry", "cellID"]
0918                 ].merge(
0919                     particle_relations,
0920                     on=["entry", "subentry"],
0921                     how="left",
0922                     validate="one_to_one",
0923                 )
0924                 sim_hit_lookup["particle_index"] = (
0925                     sim_hit_lookup["particle_index"].fillna(-1).astype(np.int64)
0926                 )
0927                 sim_particle_lookups[sim_hit_name] = sim_hit_lookup.drop_duplicates(
0928                     ["entry", "cellID"], keep="first"
0929                 )
0930
0931             sim_hit_lookup = sim_particle_lookups[sim_hit_name]
0932             if sim_hit_lookup.empty:
0933                 continue
0934
0935             rec_hit_mask = df_hits["hits_colName"].eq(rec_hit_name)
0936             rec_hits_to_match = df_hits.loc[
0937                 rec_hit_mask, ["entry", "cellID"]
0938             ].copy()
0939             rec_hits_to_match["_df_hits_index"] = rec_hits_to_match.index
0940             matched_particles = rec_hits_to_match.merge(
0941                 sim_hit_lookup[["entry", "cellID", "particle_index"]],
0942                 on=["entry", "cellID"],
0943                 how="left",
0944             )
0945             df_hits.loc[
0946                 matched_particles["_df_hits_index"].to_numpy(),
0947                 "particle_index",
0948             ] = matched_particles["particle_index"].fillna(-1).to_numpy()
0949
0950     # A trajectory relation stores the source measurement collection ID and
0951     # source object index. Match both values; measurement indices are only
0952     # unique within a collection and event.
0953     trajectory_relations, _ = get_branch_relation(
0954         tree,
0955         branch_name=traj_name,
0956         relation_name=measurement_name,
0957         relation_variables=[],
0958     )
0959
0960     # get_branch_relation preserves a trajectory with an empty relation as one
0961     # row whose relation index and collection ID are NaN. Such a row represents
0962     # no hit (for example, a trajectory with nOutliers == 0), so do not label it
0963     # as a noise hit in the returned hit table. Keep non-empty but invalid PODIO
0964     # references; those are real relation entries and remain particle_index=-1.
0965     relation_index = f"{measurement_name}_index"
0966     relation_collection = f"{measurement_name}_colID"
0967     trajectory_relations = trajectory_relations[
0968         trajectory_relations[relation_index].notna()
0969         & trajectory_relations[relation_collection].notna()
0970     ].reset_index(drop=True)
0971
0972     if df_hits.empty:
0973         traj_hits = trajectory_relations.copy()
0974         traj_hits["particle_index"] = -1
0975     else:
0976         hit_lookup = df_hits.drop_duplicates(
0977             ["entry", "measurement_colID", "subentry"],
0978             keep="first",
0979         ).rename(columns={"subentry": "_measurement_index"})
0980
0981         traj_hits = trajectory_relations.merge(
0982             hit_lookup,
0983             left_on=[
0984                 "entry",
0985                 f"{measurement_name}_colID",
0986                 f"{measurement_name}_index",
0987             ],
0988             right_on=["entry", "measurement_colID", "_measurement_index"],
0989             how="left",
0990             sort=False,
0991         )
0992         traj_hits.drop(columns=["_measurement_index"], inplace=True)
0993         traj_hits["particle_index"] = (
0994             traj_hits["particle_index"].fillna(-1).astype(np.int64)
0995         )
0996
0997     if {"position.x", "position.y"}.issubset(traj_hits.columns):
0998         traj_hits["position.r"] = np.sqrt(
0999             traj_hits["position.x"]**2 + traj_hits["position.y"]**2
1000         )
1001         traj_hits["position.phi"] = np.arctan2(
1002             traj_hits["position.y"], traj_hits["position.x"]
1003         )
1004     else:
1005         traj_hits["position.r"] = np.nan
1006         traj_hits["position.phi"] = np.nan
1007
1008     return traj_hits
1009
1010 from lmfit.models import GaussianModel
1011
1012 def gaussian(x, amplitude, mean, std):
1013     return amplitude * np.exp(-0.5 * ((x - mean) / std) ** 2) / (std * np.sqrt(2 * np.pi))
1014
1015 def hist_gaus(
1016     data, ax,
1017     bins=100, klog=False, header=None,
1018     clip=3.0, max_iters=5, min_points=50,
1019     verbose=False,
1020 ):
1021     data = np.asarray(data, dtype=float)
1022     data = data[np.isfinite(data)]
1023     if data.size < min_points:
1024         if verbose:
1025             print('hist_gaus: not enough finite points')
1026         return np.nan, np.nan, np.nan
1027
1028     center = np.median(data)
1029     mad = np.median(np.abs(data - center))
1030     scale = 1.4826 * mad if mad > 0 else np.std(data)
1031     if not np.isfinite(scale) or scale <= 0:
1032         if verbose:
1033             print('hist_gaus: invalid initial scale')
1034         return np.nan, np.nan, np.nan
1035
1036     for _ in range(max_iters):
1037         mask = np.abs(data - center) <= clip * scale
1038         clipped = data[mask]
1039         if clipped.size < min_points:
1040             break
1041         new_center = np.median(clipped)
1042         mad = np.median(np.abs(clipped - new_center))
1043         new_scale = 1.4826 * mad if mad > 0 else np.std(clipped)
1044         if not np.isfinite(new_scale) or new_scale <= 0:
1045             break
1046         if np.isclose(new_center, center) and np.isclose(new_scale, scale):
1047             center, scale = new_center, new_scale
1048             break
1049         center, scale = new_center, new_scale
1050
1051     lo, hi = center - clip * scale, center + clip * scale
1052     if lo == hi:
1053         if verbose:
1054             print('hist_gaus: degenerate range')
1055         return np.nan, np.nan, np.nan
1056
1057     counts, edges = np.histogram(data, bins=bins, range=(lo, hi))
1058     mid = 0.5 * (edges[:-1] + edges[1:])
1059
1060     if ax is not None:
1061         ax.hist(data, bins=bins, range=(lo, hi), histtype='stepfilled', alpha=0.3)
1062
1063     model = GaussianModel()
1064     params = model.make_params(center=center, sigma=scale, amplitude=np.max(counts))
1065     try:
1066         result = model.fit(counts, params, x=mid)
1067     except Exception as exc:
1068         if verbose:
1069             print('hist_gaus: fit failed', exc)
1070         return np.nan, np.nan, np.nan
1071
1072     sigma = float(result.params['sigma'])
1073     sigma_err = result.params['sigma'].stderr
1074     if sigma_err is None or not np.isfinite(sigma) or sigma <= 0:
1075         if verbose:
1076             print('hist_gaus: invalid fit result')
1077         return np.nan, np.nan, np.nan
1078
1079     mean = float(result.params['center'])
1080     sigma_err = float(sigma_err)
1081     ampl = float(result.params['amplitude'])
1082     peak = ampl / (sigma * np.sqrt(2 * np.pi))
1083
1084     if ax is not None:
1085         ax.plot(mid, gaussian(mid, ampl, mean, sigma), 'r-')
1086         if header:
1087             ax.set_title(header)
1088         ax.set_xlabel('value')
1089         ax.set_ylabel('entries')
1090         if klog:
1091             ax.set_yscale('log')
1092         else:
1093             ymax = max(np.max(counts), peak)
1094             ax.set_ylim(0, ymax / 0.7)
1095
1096     return mean, sigma, sigma_err
1097
1098
1099 __all__ = [
1100     "configure_analysis_environment",
1101     "ak_flat",
1102     "ak_df",
1103     "ak_hist",
1104     "ak_filter",
1105     "get_pdg_info",
1106     "get_geoID",
1107     "theta2eta",
1108     "select_string",
1109     "read_ur",
1110     "get_col_table",
1111     "get_branch_ak",
1112     "get_branch_df",
1113     "get_part",
1114     "get_params",
1115     "get_branches",
1116     "get_collections",
1117     "get_relation",
1118     "get_branch_relation",
1119     "get_traj_relations",
1120     "get_traj_hits_particle",
1121     "deg2rad",
1122     "status_to_source",
1123     "geo_mask_values",
1124     "hist_gaus"
1125 ]