Tracking/track2particle/epic_analysis_podio.py

0001 '''
0002     Podio-specific helpers for epic tracking analysis.
0003     Keep this module optional and import only in environments with podio.
0004     Shujie Li, Aug 2025
0005 '''
0006
0007 import numpy as np
0008 import pandas as pd
0009 from matplotlib import pyplot as plt
0010 import fnmatch
0011
0012 from epic_analysis_base import (
0013     TRACK_HIT_COUNT_MIN_MIN,
0014     TRACK_HIT_COUNT_MIN,
0015     TRACK_MOM_MIN,
0016     TRACK_HIT_FRACTION_MIN,
0017     TRACK_HIT_COUNT_GHOST_MAX,
0018     VERTEX_CUT_R_MAX,
0019     VERTEX_CUT_Z_MAX,
0020     status_to_source,
0021 )
0022
0023 def read_podio(fname, s3_dir="",tname="events"):
0024     """Read ROOT file with podio. Does NOT work for the metadata tree """
0025     # from podio import root_io
0026     import sys
0027     if 'podio' not in sys.modules:
0028         print("Loading podio ROOT IO reader(this will take ~2 minutes)...")
0029
0030     from podio.root_io import Reader  # More specific
0031
0032     server = 'root://dtn-eic.jlab.org//volatile/eic/'
0033     if len(s3_dir) > 1:
0034         fname = server + s3_dir + fname
0035     reader = Reader(fname)
0036     tree = reader.get("events")
0037     print(f"read_podio: read {fname}:{tname}. {len(tree)} events in total")
0038     return tree
0039
0040 def show_getter_podio(collection):
0041     """Display basic info about collection or single object"""
0042     type_check = check_type(collection)
0043
0044     if type_check.is_iterable:
0045         print(f"Number of objects: {len(collection)}")
0046         sample_obj = collection[0] if len(collection) > 0 else None
0047     else:
0048         print("Single object")
0049         sample_obj = collection
0050
0051     if sample_obj:
0052         sample_type = type(sample_obj).__name__
0053         print(f"Object type: {sample_type}")
0054         getters = [m for m in dir(sample_obj) if m.startswith('get') and not m.startswith('_')]
0055         print(f"Available getter methods ({len(getters)}):")
0056         for getter in getters:
0057             print(f"  {getter}")
0058
0059         # Pretty print version
0060         scalars = {m: getattr(sample_obj, m)() for m in dir(sample_obj)
0061                 if m.startswith('get') and not m.startswith('_') and callable(getattr(sample_obj, m))
0062                 and isinstance(getattr(sample_obj, m)(), (int, float, bool, str))}
0063         for k, v in scalars.items(): print(f"{k}: {v}")
0064
0065 def show_collections_podio(event, pattern=None):
0066     """Show collections matching pattern (case-insensitive), or all if no pattern"""
0067     all_collections = event.getAvailableCollections()
0068     if pattern:
0069         collections = [name for name in all_collections
0070                       if fnmatch.fnmatch(name.lower(), pattern.lower())]
0071     else:
0072         collections = all_collections
0073     # print(f"Collections ({len(collections)}):")
0074     # for name in collections:
0075     #     print(f"  {name}")
0076     return collections
0077
0078 def get_collection_member_podio(podio_collection, member_name):
0079     """
0080     Access podio object members in uproot/awkward style
0081
0082     Args:
0083         podio_collection: The podio branch (collection or single object)
0084         member_name: Member name (like "nHoles", "chi2")
0085     """
0086     type_check = check_type(podio_collection)
0087     if type_check.is_empty:
0088         raise ValueError("Collection is empty")
0089
0090     # Helper function to process a single object
0091     def process_single_object(obj):
0092         # Common getter patterns in podio
0093         possible_getters = [
0094             f"get{member_name}",
0095             f"get{member_name.capitalize()}",
0096             f"getN{member_name.capitalize()}",  # for count-like members
0097             member_name
0098         ]
0099
0100         for getter_name in possible_getters:
0101             if hasattr(obj, getter_name):
0102                 attr = getattr(obj, getter_name)
0103                 if callable(attr):
0104                     return attr()
0105                 else:
0106                     return attr
0107
0108         # If not found, show available members
0109         all_methods = [method for method in dir(obj) if not method.startswith('_')]
0110         getters = [method for method in all_methods if method.startswith('get')]
0111         other_attrs = [attr for attr in all_methods if not attr.startswith('get') and not callable(getattr(obj, attr, None))]
0112
0113         error_msg = f"Cannot find member '{member_name}' in {type(obj).__name__}\n"
0114         error_msg += f"Available getter methods:\n"
0115         for getter in getters:
0116             error_msg += f"  {getter}\n"
0117         if other_attrs:
0118             error_msg += f"Available attributes:\n"
0119             for attr in other_attrs:
0120                 error_msg += f"  {attr}\n"
0121
0122         raise AttributeError(error_msg)
0123
0124     if type_check.is_iterable:
0125         # Process collection - return list of results
0126         results = []
0127         for obj in podio_collection:
0128             results.append(process_single_object(obj))
0129         return results
0130     else:
0131         # Process single object - return single result
0132         return process_single_object(podio_collection)
0133
0134 class PodioCollectionWrapper:
0135     def __init__(self, podio_collection):
0136         self.collection = podio_collection
0137         self.type_check = check_type(podio_collection)
0138
0139     def __getitem__(self, member_name):
0140         return get_collection_member_podio(self.collection, member_name)
0141
0142     def __len__(self):
0143         return len(self.collection) if self.type_check.is_iterable else 1
0144
0145     def __iter__(self):
0146         return iter(self.collection) if self.type_check.is_iterable else iter([self.collection])
0147
0148 def check_type(obj):
0149     """Check object type with comprehensive categorization"""
0150     def get_value_category(value):
0151         # Check for numbers first
0152         if isinstance(value, (int, float, complex, bool)):
0153             return 'number'
0154         # Check for numpy arrays
0155         elif isinstance(value, np.ndarray):
0156             if value.ndim == 1:
0157                 return 'array'
0158             else:
0159                 return 'nested_array'  # 2D, 3D arrays etc.
0160         # Check for Python lists/tuples
0161         elif isinstance(value, (list, tuple)):
0162             if len(value) > 0:
0163                 # Check if it's a list of arrays (nested structure)
0164                 first_item = value[0]
0165                 if isinstance(first_item, (np.ndarray, list, tuple)):
0166                     return 'nested_array'
0167                 else:
0168                     return 'list'
0169             else:
0170                 return 'list'
0171         # Check for podio RelationRange or other iterable collections
0172         elif 'RelationRange' in str(type(value)) or (hasattr(value, '__len__') and hasattr(value, '__iter__') and not isinstance(value, str)):
0173             return 'range'
0174         # Everything else is an object
0175         else:
0176             return 'object'
0177
0178     # Create a simple container
0179     class CategoryValue:
0180         def __init__(self, val):
0181             self.value = val
0182             cat = get_value_category(val)
0183             self.category = cat
0184             self.size        = getattr(val, '__len__', lambda: 1)()
0185             self.is_empty    = getattr(val, '__len__', lambda: 1)() == 0
0186
0187             self.is_range  = (cat == "range")
0188             self.is_number = (cat == "number")
0189             self.is_object = (cat == "object")
0190             self.is_array  = (cat == "array")
0191             self.is_nested_array = (cat == "nested_array")
0192             self.is_list   = (cat == "list")
0193             # Convenience groupings
0194             self.is_iterable = cat in ["range", "array", "nested_array", "list"]
0195             self.is_numpy    = cat in ["array", "nested_array"]
0196             self.is_simple   = cat in ["number", "list", "array"]
0197
0198     return CategoryValue(obj)
0199
0200
0201 def build_rawhit_lookup(obj_list):
0202     """Build fast lookup using cellID + time + charge as unique key"""
0203     return {(obj.getCellID(), obj.getTimeStamp(), obj.getCharge()): i
0204             for i, obj in enumerate(obj_list)}
0205
0206 def get_rawhit_index(lookup,obj):
0207     key = (obj.getCellID(), obj.getTimeStamp(), obj.getCharge())
0208     return lookup.get(key, -1)
0209
0210 def is_valid_podio_object(obj):
0211     """Return False for a null or unresolved PODIO relation."""
0212     if obj is None:
0213         return False
0214     try:
0215         object_id = obj.id()
0216     except (AttributeError, ReferenceError, RuntimeError):
0217         return False
0218     return object_id.index >= 0 and object_id.collectionID != 0xFFFFFFFF
0219
0220 def build_obj_lookup(obj_list):
0221     """Build lookup that handles multiple objects with same collectionID+index"""
0222     lookup = {}
0223     for i, obj in enumerate(obj_list):
0224         if not is_valid_podio_object(obj):
0225             continue
0226         key = (obj.id().collectionID, obj.id().index)
0227         if key not in lookup:
0228             lookup[key] = []
0229         lookup[key].append(i)
0230     return lookup
0231
0232 def get_obj_indices(lookup, obj):
0233     """Return list of all indices matching the object's key"""
0234     if not is_valid_podio_object(obj):
0235         return []
0236     key = (obj.id().collectionID, obj.id().index)
0237     return lookup.get(key, [])
0238
0239 def get_traj_hits(event,bname="CentralCKFTrajectories",kcombine=0):
0240     ## ALL raw (one) --> sim hit (many, with weight) associations for a given event. Noise hits won't have association.
0241     ## ----------------
0242     # traj-based info
0243     ## ----------------
0244     ## keep track of the rawhit(simhit) index in the association, which includes all central tracker hits in one collection.
0245     asso=event.get("CentralTrackingRawHitAssociations")
0246     asso_raw=PodioCollectionWrapper(asso)["RawHit"]
0247     lookup = build_obj_lookup(asso_raw)
0248     br    = event.get(bname) ## for one event, can have multiple subentries
0249     vname = "Measurements_deprecated"
0250     ltraj=[]
0251     lhit=[]
0252     lweight=[]
0253     # lpos =[]
0254     lpart=[]
0255     lrec_hit_col=[]
0256     lrec_hit_id=[]
0257     lmeasurement_col=[]
0258     lmeasurement_id=[]
0259     lhit_in_measurement=[]
0260     def add_invalid_hit(traj_id, measurement_col, measurement_id, hit_number,
0261                         rec_col=-1, rec_id=-1, association_id=-1):
0262         """Keep an invalid or noise hit in the trajectory accounting."""
0263         ltraj.append(traj_id)
0264         lhit.append(association_id)
0265         lpart.append(-1)
0266         lrec_hit_col.append(rec_col)
0267         lrec_hit_id.append(rec_id)
0268         lmeasurement_col.append(measurement_col)
0269         lmeasurement_id.append(measurement_id)
0270         lhit_in_measurement.append(hit_number)
0271
0272     ## for each traj--> each measurement --> rec hit --> raw hit --> match raw with sim by association index --> particle
0273     for ii,traj in enumerate(br):
0274         measurements  = PodioCollectionWrapper(traj)[vname]
0275         for measurement in measurements:
0276             if not is_valid_podio_object(measurement):
0277                 add_invalid_hit(ii, -1, -1, -1)
0278                 continue
0279
0280             measurement_col = measurement.id().collectionID
0281             measurement_id  = measurement.id().index
0282             hits = measurement.getHits()
0283             if len(hits) == 0:
0284                 add_invalid_hit(ii, measurement_col, measurement_id, -1)
0285                 continue
0286
0287             # A Measurement2D can contain more than one constituent hit. In
0288             # particular, TOF clusters commonly contain two shared rec hits.
0289             for hit_number, hit in enumerate(hits):
0290                 if not is_valid_podio_object(hit):
0291                     add_invalid_hit(
0292                         ii, measurement_col, measurement_id, hit_number
0293                     )
0294                     continue
0295
0296                 rec_col = hit.id().collectionID
0297                 rec_id  = hit.id().index
0298                 raw = hit.getRawHit() # rec2raw is one-to-one
0299                 if not is_valid_podio_object(raw):
0300                     add_invalid_hit(
0301                         ii, measurement_col, measurement_id, hit_number,
0302                         rec_col, rec_id
0303                     )
0304                     continue
0305
0306                 # One raw hit can occur more than once when it is associated
0307                 # with several simulated hits.
0308                 indx = get_obj_indices(lookup, raw)
0309                 if len(indx) == 0:
0310                     # A valid rec hit without a truth association is noise.
0311                     add_invalid_hit(
0312                         ii, measurement_col, measurement_id, hit_number,
0313                         rec_col, rec_id
0314                     )
0315                     continue
0316
0317                 for ind in indx:
0318                     sim = asso[ind].getSimHit()
0319                     if not is_valid_podio_object(sim):
0320                         add_invalid_hit(
0321                             ii, measurement_col, measurement_id, hit_number,
0322                             rec_col, rec_id, ind
0323                         )
0324                         continue
0325
0326                     part = sim.getParticle()
0327                     if not is_valid_podio_object(part):
0328                         add_invalid_hit(
0329                             ii, measurement_col, measurement_id, hit_number,
0330                             rec_col, rec_id, ind
0331                         )
0332                         continue
0333
0334                     ltraj.append(ii)
0335                     lhit.append(ind)
0336                     lpart.append(part.id().index)
0337                     lrec_hit_col.append(rec_col)
0338                     lrec_hit_id.append(rec_id)
0339                     lmeasurement_col.append(measurement_col)
0340                     lmeasurement_id.append(measurement_id)
0341                     lhit_in_measurement.append(hit_number)
0342     traj_hits=pd.DataFrame({
0343         "traj_id": ltraj,
0344         "part_id": lpart,
0345         "asso_hit": lhit,
0346         "measurement_col": lmeasurement_col,
0347         "measurement_id": lmeasurement_id,
0348         "hit_in_measurement": lhit_in_measurement,
0349         "rec_hit_col": lrec_hit_col,
0350         "rec_hit_id": lrec_hit_id,
0351     })
0352     #"position":lpos, "hit_weight":lweight})
0353
0354     ## FIXME: check for overlapped tracks (for now ambiguity solver config won't allow sharing hits)
0355     # reoccur_hit=traj_hits.groupby('asso_hit').filter(lambda group: len(group) > 1)
0356     # for row in reoccur_hit.itertuples():
0357     #     print(f'WARNING: duplihits detected:', row)
0358
0359     ## if one rec hit is associated to multiple sim hits, but all sim hits go to the same particle, then only keep one entry.
0360     ## FIXME: not sure if this will work if we allow overlap traj hit.
0361     traj_hits = traj_hits.drop_duplicates(
0362         subset=[
0363             "part_id", "traj_id", "measurement_col", "measurement_id",
0364             "hit_in_measurement", "rec_hit_col", "rec_hit_id",
0365         ],
0366         keep="first",
0367     )
0368     traj_hits['weight'] = (traj_hits.groupby(['traj_id', 'part_id']).transform('size') /
0369                         traj_hits.groupby('traj_id').transform('size'))
0370
0371     if kcombine:
0372         traj_hits = traj_hits.groupby(['traj_id', 'part_id'], as_index=False).agg({
0373             'asso_hit': list,
0374             'weight': 'first'  # Keep the weight value (should be same for same traj+particle)
0375         })
0376
0377     return traj_hits
0378
0379 def get_part_hits(event, traj_hits, ksignal=0):
0380     ## ----------------
0381     # particle-based info
0382     ## ----------------
0383     ltraj_id=[]
0384     lpart_id=[]
0385     lsimhit=[]
0386     lgenID=[]
0387     lraw_hit_col=[]
0388     lraw_hit_id=[]
0389     # lpart=[]
0390     # lposition=[]
0391
0392     ## for fast lookup (instead of traj_hits.asso_hit==ii). Assume no shared hits
0393     # Invalid or noise hits use asso_hit=-1 and intentionally have no entry
0394     # in the truth-association lookup.
0395     traj_base = traj_hits[traj_hits["asso_hit"] >= 0].set_index("asso_hit")
0396     if not traj_base.index.is_unique:
0397         raise ValueError("get_part_hits: duplicate asso_hit in traj_hits")
0398     traj_map = traj_base['traj_id']
0399     ## for each simhit (that is converted to rec hit-->measurement candidate), find related particle and traj
0400     asso=event.get("CentralTrackingRawHitAssociations")
0401     for ii,association in enumerate(asso):
0402         sim  = association.getSimHit()
0403         if not is_valid_podio_object(sim):
0404             continue
0405         part = sim.getParticle()
0406         if not is_valid_podio_object(part):
0407             continue
0408         raw  = association.getRawHit()
0409         if not is_valid_podio_object(raw):
0410             continue
0411         # cond_vertex   =  (np.sqrt(part.getVertex().x**2+part.getVertex().y**2)<1 )&(abs(part.getVertex().z)<100)
0412         # cond = cond_vertex
0413         status = part.getGeneratorStatus()
0414         cond = (status in (1, 2)) if ksignal == 1 else True
0415         if cond:
0416             lsimhit.append(ii) ## as before, use unique index (and unique sim hit) from the association.
0417             # lpart.append(part)
0418             lpart_id.append(part.id().index)
0419             lgenID.append(status)
0420             lraw_hit_col.append(raw.id().collectionID)
0421             lraw_hit_id.append(raw.id().index)
0422             # lposition.append(sim.getPosition())
0423
0424             ## find which trajectory used this hit
0425             ltraj_id.append(int(traj_map.get(ii, -1)))
0426     part_hits = pd.DataFrame({"part_id":lpart_id, "part_status":lgenID, "asso_hit":lsimhit, "traj_id":ltraj_id,
0427                             "raw_hit_col": lraw_hit_col, "raw_hit_id": lraw_hit_id})#, "position":lposition, "particle": lpart})
0428     part_hits = part_hits.drop_duplicates(subset=["part_id", "traj_id", "raw_hit_col", "raw_hit_id"],  keep="first")
0429
0430     # df = primary_hits[primary_hits.groupby('particle')['particle'].transform('count') >= 3]
0431     return part_hits
0432
0433 def get_traj_purity(traj_hits):
0434     '''
0435     traj_hits is the output from get_traj_hits()
0436     returns trajectory and source, max fraction=purity
0437     '''
0438     grouped = traj_hits.groupby(['traj_id'])
0439     # Analyze each group
0440     result = grouped['part_id'].agg([
0441         ('total_count', 'count'),
0442         ('unique_source', lambda x: x.nunique()),
0443         ('most_common_source', lambda x: x.value_counts().idxmax()),
0444         ('max_count', lambda x: x.value_counts().max())
0445     ])
0446
0447     # Calculate derived columns
0448     result['max_fraction'] = result['max_count'] / result['total_count']
0449     # result['all_same'] = result['unique_source'] == 1
0450     return result.copy().reset_index()
0451
0452
0453 def plot_part_traj_flow(df, params=None, mcpart=None):
0454     """Create alluvial-style diagram showing particle-trajectory flows"""
0455     # dict_part = dict(zip(df["part_id"], df["particle"]))
0456
0457     # Separate used and unused hits
0458     df_used = df[df['traj_id'] != -1]  # Only hits used in trajectories
0459     df_all = df  # All hits including unused
0460
0461     # Calculate statistics
0462     particle_totals = df_all.groupby('part_id')['asso_hit'].apply(len)    # Total hits per particle
0463     particle_used   = df_used.groupby('part_id')['asso_hit'].apply(len)   # Used hits per particle
0464     traj_totals     = df_used.groupby('traj_id')['asso_hit'].apply(len)   # Total hits per trajectory
0465
0466     # Only include particles that have some used hits (have flows)
0467     particles_with_flows = df_used['part_id'].unique()
0468     trajectories = sorted(df_used['traj_id'].unique())  # Only trajectories with used hits
0469     len1 = len(particles_with_flows)
0470     len2 = len(trajectories)
0471     # Prepare flow data (only for used hits)
0472     flows = df_used.groupby(['part_id', 'traj_id'])['asso_hit'].apply(len).reset_index()
0473     flows.columns = ['part_id', 'traj_id', 'hits']
0474
0475     fig, ax = plt.subplots(figsize=(12, 8))
0476
0477     # Position particles on left, trajectories on right
0478     particle_y = {p: i for i, p in enumerate(sorted(particles_with_flows))}
0479     traj_y     = {t: i for i, t in enumerate(trajectories)}
0480
0481     # Draw flows as curved lines
0482     for _, row in flows.iterrows():
0483         particle = row['part_id']
0484         traj     = row['traj_id']
0485         hits     = row['hits']
0486
0487         # Start and end points
0488         x1, y1 = 0, len1 - particle_y[particle]
0489         x2, y2 = 1, len1 - traj_y[traj]
0490
0491         # Create curved line
0492         x_curve = [x1, 0.5, x2]
0493         y_curve = [y1, (y1 + y2) / 2, y2]
0494
0495         # Line thickness proportional to hits
0496         linewidth = max(1, hits / flows['hits'].max() * 10)
0497
0498         ax.plot(x_curve, y_curve, linewidth=linewidth, alpha=0.6)
0499
0500     # Add particle labels (only for particles with flows)
0501     ax.text(-0.25, len1+1,  f'Particle ID:  (used/total hits)',  ha='left', va='center')
0502     ax.text(0.9,   len1+1,  f'Trajectory ID: (nMeasurements)',  ha='left', va='center')
0503
0504     for i, p in enumerate(sorted(particles_with_flows)):
0505         total_hits = particle_totals[p]
0506         used_hits = particle_used[p]
0507         color = 'k'
0508         if total_hits<TRACK_HIT_COUNT_MIN:
0509             color="grey"
0510         if mcpart is not None:
0511             pp=mcpart.iloc[int(p)]
0512             if abs(pp["vertex_r"])>VERTEX_CUT_R_MAX or abs(pp["vertex.z"])>VERTEX_CUT_Z_MAX or pp["mom"]<TRACK_MOM_MIN:
0513                 color='grey'
0514             ax.text(-0.03, len1-i,  f'#{p}: {status_to_source[pp.generatorStatus]} ({used_hits}/{total_hits})',
0515                 ha='right', va='center',color=color)
0516         else:
0517             ax.text(-0.03, len1-i,  f'#{p}:   ({used_hits}/{total_hits})',
0518                 ha='right', va='center',color=color)
0519
0520     # Add trajectory labels
0521     for i, t in enumerate(trajectories):
0522         total_traj_hits = traj_totals[t]
0523         color="k"
0524         if total_traj_hits<TRACK_HIT_COUNT_MIN:
0525             color="grey"
0526         if params is not None:
0527             p = params.iloc[int(t)]
0528             if abs(p["loc.a"])>VERTEX_CUT_R_MAX or abs(p["loc.b"])>VERTEX_CUT_Z_MAX or p["mom"]<TRACK_MOM_MIN:
0529                 color='grey'
0530                 # print(t, ["loc.a"], p["loc.b"])
0531         ax.text(1.03, len1-i, f'#{int(t)} ({total_traj_hits})',
0532                 ha='left', va='center', color=color)
0533
0534     ax.set_xlim(-0.2, 1.2)
0535     ax.set_ylim(-0.5, max(len(particles_with_flows), len(trajectories)) +3)
0536     ax.set_title('Particle to Trajectory Flow (Particles with Used Hits Only)')
0537     ax.axis('off')
0538
0539     return plt
0540
0541
0542 def get_part_traj_counts(event,mcpart, ksignal=0, kverbose=0):
0543     traj_hits=get_traj_hits(event)
0544     part_hits=get_part_hits(event,traj_hits, ksignal)
0545
0546     # Prefer explicit particle-id columns; fall back to positional index for compatibility.
0547     id_col = None
0548     if "orig_subentry" in mcpart.columns:
0549         id_col = "orig_subentry"
0550     elif "subentry" in mcpart.columns:
0551         id_col = "subentry"
0552
0553     ## -----------Find good track-----------
0554     ## Get counts per (traj_id, part_id) and total per traj_id
0555     traj_counts = get_traj_purity(traj_hits)
0556     ## get generator status
0557     if id_col is not None:
0558         status_map = (
0559             mcpart.drop_duplicates(subset=[id_col], keep="first")
0560             .set_index(id_col)["generatorStatus"]
0561         )
0562         traj_counts["part_status"] = (
0563             traj_counts["most_common_source"].map(status_map).fillna(-1).astype(int)
0564         )
0565     else:
0566         traj_counts["part_status"]=traj_counts["most_common_source"].apply(lambda x: mcpart.iloc[x].generatorStatus)
0567
0568     ## -----------Find good particles-----------
0569     ## track hit cut
0570     part_counts = part_hits.groupby("part_id").size()
0571     part_counts = part_counts[part_counts>=TRACK_HIT_COUNT_MIN_MIN]
0572     if id_col is not None:
0573         mcpart_hits = mcpart[mcpart[id_col].isin(part_counts.index)].copy()
0574         mcpart_hits["hit_counts"] = mcpart_hits[id_col].map(part_counts).astype(int)
0575     else:
0576         mcpart_hits = mcpart.iloc[part_counts.index].copy()
0577         mcpart_hits["hit_counts"] = part_counts.values
0578     # traj_counts["part_status"]=traj_counts["most_common_source"].apply(lambda x: mcpart.iloc[x].generatorStatus)
0579     ## only do event-by-event quality check when required. Otherwise return the dataframe for further analysis
0580     if kverbose:
0581         good_part_id   = (part_hits.groupby("part_id").size()>=TRACK_HIT_COUNT_MIN)
0582         part_hits_good = part_hits[part_hits["part_id"].isin(good_part_id[good_part_id].index)][["part_id","part_status"]].drop_duplicates()
0583         if id_col is not None:
0584             good_mcpart = mcpart[mcpart[id_col].isin(part_hits_good.part_id.unique())].copy()
0585         else:
0586             good_mcpart = mcpart.iloc[part_hits_good.part_id.unique()]
0587         ## vertex and momentum cut
0588         cond_vertex = (abs(good_mcpart.vertex_r)<VERTEX_CUT_R_MAX)&(abs(good_mcpart["vertex.z"])<VERTEX_CUT_Z_MAX)
0589         cond_mom    = (good_mcpart.mom>TRACK_MOM_MIN)
0590         good_mcpart = good_mcpart[cond_vertex&cond_mom]
0591         ## signal or background
0592         cond_sig    = (good_mcpart.generatorStatus==1)|(good_mcpart.generatorStatus==2)
0593         good_mcpart_sig  =good_mcpart[cond_sig]
0594         good_mcpart_other=good_mcpart[(~cond_sig)]
0595         print("Number of good particles (signal, other):",len(good_mcpart_sig), len(good_mcpart_other))
0596
0597
0598     if kverbose:
0599         traj_counts["traj_status"]=0
0600         traj_counts.loc[(traj_counts.max_fraction<TRACK_HIT_FRACTION_MIN) | (traj_counts.max_count<=TRACK_HIT_COUNT_GHOST_MAX),"traj_status"]=-1 ## ghost status=-1
0601         traj_counts.loc[(traj_counts.traj_status>-1)&(traj_counts['total_count']>=TRACK_HIT_COUNT_MIN),'traj_status']=1
0602
0603         ghost_traj_id = traj_counts[traj_counts.traj_status==-1].index.tolist()
0604         print("list of ghost track id:", ghost_traj_id)
0605         good_traj=traj_counts[traj_counts.traj_status==1].copy()
0606         good_traj_sig=good_traj[(good_traj.part_status==1)|(good_traj.part_status==2)]
0607         ntraj_sig   = len(good_traj_sig)
0608         ntraj_other = len(good_traj)-ntraj_sig
0609         print("Number of good tracks from sig/others:",ntraj_sig, ntraj_other)
0610     ## hit purity for good tracks
0611     ## Percentage of track hits from a single source
0612         purity_hit=good_traj.max_fraction.to_list()
0613         print("Hit purity:" , purity_hit)
0614     ##-----------Track to particle efficiency----------------
0615     # fraction of good signal particles that are linked to some good track
0616         if id_col is not None:
0617             good_part_sig_id = good_mcpart_sig[id_col].tolist()
0618         else:
0619             good_part_sig_id = good_mcpart_sig.index.get_level_values('subentry').tolist()
0620         good_traj_id = good_traj.most_common_source.unique()
0621         good_traj_sig_id = good_traj_sig.most_common_source.unique()
0622         common = list(set(good_part_sig_id) & set(good_traj_id))
0623         print(set(good_part_sig_id),  set(good_traj_sig_id), common)
0624         ## 50% of the total particle hits go to that traj <---- skip this for now for a looser check
0625         # good_traj['fract'] = good_traj['max_count'] / good_traj['most_common_source'].map(part_hits['part_id'].value_counts())
0626         # good_traj=good_traj.fract>=0.5
0627     return mcpart_hits, traj_counts
0628
0629
0630
0631 __all__ = [
0632     "read_podio",
0633     "show_getter_podio",
0634     "show_collections_podio",
0635     "get_collection_member_podio",
0636     "PodioCollectionWrapper",
0637     "check_type",
0638     "build_rawhit_lookup",
0639     "get_rawhit_index",
0640     "is_valid_podio_object",
0641     "build_obj_lookup",
0642     "get_obj_indices",
0643     "get_traj_hits",
0644     "get_part_hits",
0645     "get_traj_purity",
0646     "plot_part_traj_flow",
0647     "get_part_traj_counts",
0648 ]