Tracking/svt_rate_per_source/svt_sources.py

0001 """SVT geometry selections, PODIO reading, and rec-hit source assignment.
0002
0003 All coordinates are mm. These windows reproduce the parent workflow's SVT
0004 selections; they are analysis cuts, not a geometry or sensor-layout service.
0005 """
0006
0007 from collections import Counter, defaultdict
0008
0009 import awkward as ak
0010 import numpy as np
0011
0012
0013 STATUS_TO_SOURCE = {
0014     1: "DIS", 2: "DIS",
0015     2001: "SR", 2002: "SR",
0016     3001: "Bremstrahlung", 3002: "Bremstrahlung",
0017     4001: "Coulomb", 4002: "Coulomb",
0018     5001: "Touschek", 5002: "Touschek",
0019     6001: "Proton beam gas", 6002: "Proton beam gas",
0020 }
0021 BACKGROUND_COMPONENTS = {"Bremstrahlung", "Coulomb", "Touschek", "Proton beam gas"}
0022 SOURCE_ORDER = ["ALL", "DIS", "SR", "other background",
0023                 "Bremstrahlung", "Coulomb", "Touschek", "Proton beam gas",
0024                 "other", "unmatched", "invalid_relation"]
0025
0026 # (surface name, rec-hit collection, barrel?, radius or z window in mm).
0027 SURFACES = [
0028     ("L0", "SiBarrelVertexRecHits", True, (30, 42)),
0029     ("L1", "SiBarrelVertexRecHits", True, (46, 60)),
0030     ("L2", "SiBarrelVertexRecHits", True, (115, 130)),
0031     ("L3", "SiBarrelTrackerRecHits", True, (250, 290)),
0032     ("L4", "SiBarrelTrackerRecHits", True, (400, 450)),
0033     ("E-Si Disk 4", "SiEndcapTrackerRecHits", False, (-1055, -1000)),
0034     ("E-Si Disk 3", "SiEndcapTrackerRecHits", False, (-860, -840)),
0035     ("E-Si Disk 2", "SiEndcapTrackerRecHits", False, (-660, -640)),
0036     ("E-Si Disk 1", "SiEndcapTrackerRecHits", False, (-460, -440)),
0037     ("E-Si Disk 0", "SiEndcapTrackerRecHits", False, (-260, -240)),
0038     ("H-Si Disk 0", "SiEndcapTrackerRecHits", False, (240, 260)),
0039     ("H-Si Disk 1", "SiEndcapTrackerRecHits", False, (440, 460)),
0040     ("H-Si Disk 2", "SiEndcapTrackerRecHits", False, (690, 710)),
0041     ("H-Si Disk 3", "SiEndcapTrackerRecHits", False, (940, 960)),
0042     ("H-Si Disk 4", "SiEndcapTrackerRecHits", False, (1150, 1250)),
0043 ]
0044
0045 # Explicit relation paths keep the EDM convention auditable.
0046 RELATIONS = [
0047     {
0048         "rec": "SiEndcapTrackerRecHits",
0049         "rec_raw": "_SiEndcapTrackerRecHits_rawHit",
0050         "association": "SiEndcapTrackerRawHitAssociations",
0051         "association_raw": "_SiEndcapTrackerRawHitAssociations_rawHit",
0052         "association_sim": "_SiEndcapTrackerRawHitAssociations_simHit",
0053         "sim_particle": "_TrackerEndcapHits_particle",
0054     },
0055     {
0056         "rec": "SiBarrelVertexRecHits",
0057         "rec_raw": "_SiBarrelVertexRecHits_rawHit",
0058         "association": "SiBarrelVertexRawHitAssociations",
0059         "association_raw": "_SiBarrelVertexRawHitAssociations_rawHit",
0060         "association_sim": "_SiBarrelVertexRawHitAssociations_simHit",
0061         "sim_particle": "_VertexBarrelHits_particle",
0062     },
0063     {
0064         "rec": "SiBarrelTrackerRecHits",
0065         "rec_raw": "_SiBarrelTrackerRecHits_rawHit",
0066         "association": "SiBarrelRawHitAssociations",
0067         "association_raw": "_SiBarrelRawHitAssociations_rawHit",
0068         "association_sim": "_SiBarrelRawHitAssociations_simHit",
0069         "sim_particle": "_SiBarrelHits_particle",
0070     },
0071 ]
0072
0073
0074 def read_branch(tree, name, start, stop, fields=None):
0075     """Read one event range and shorten split PODIO names (e.g. position.x).
0076
0077     The outer workflow handles chunking. All required fields must be present;
0078     missing collections or incompatible schemas fail rather than yield zeros.
0079     """
0080     try:
0081         data = tree[name].array(library="ak", entry_start=start, entry_stop=stop)
0082     except KeyError as error:
0083         raise KeyError(f"Missing required PODIO collection/relation: {name}") from error
0084     renamed = {}
0085     for field in data.fields:
0086         short = field[len(name) + 1:] if field.startswith(name + ".") else field
0087         if "[" not in short and (fields is None or short in fields):
0088             renamed[short] = data[field]
0089     if fields is not None:
0090         missing = set(fields) - renamed.keys()
0091         if missing:
0092             raise ValueError(f"{name}: missing fields {sorted(missing)}")
0093     if not renamed:
0094         raise ValueError(f"{name}: expected a split PODIO record, got {data.type}")
0095     return ak.zip(renamed)
0096
0097
0098 def resolve_status(index, statuses, begin, end, parents):
0099     """Follow the first parent of status-0 secondaries, with a 100-link limit.
0100
0101     Returns (resolved status, traced secondary?, unresolved ancestry?). None
0102     denotes an invalid initial MC index; an unresolved secondary retains 0.
0103     Multiple parents are deliberately not combined, matching the parent code.
0104     """
0105     if index < 0 or index >= len(statuses):
0106         return None, False, False
0107     if statuses[index] != 0:
0108         return int(statuses[index]), False, False
0109     seen = set()
0110     for _ in range(100):
0111         if index in seen:
0112             break
0113         seen.add(index)
0114         first, last = int(begin[index]), int(end[index])
0115         if first < 0 or first >= last or last > len(parents):
0116             break
0117         index = int(parents[first])
0118         if index < 0 or index >= len(statuses):
0119             break
0120         if statuses[index] != 0:
0121             return int(statuses[index]), True, False
0122     return 0, True, True
0123
0124
0125 def assign_sources(rec_raw, weights, association_raw, association_sim,
0126                    sim_particle, mc_particles, mc_parents):
0127     """Assign one exclusive source to each rec hit in one event.
0128
0129     ObjectID indices are event-local within the named collections in RELATIONS.
0130     Choose the largest association weight; ties choose the first row. These
0131     assumptions follow the original analysis. Collection IDs are not remapped.
0132     Invalid selected relations remain visible instead of trying another source.
0133     """
0134     raw_indices = ak.to_numpy(association_raw["index"])
0135     sim_indices = ak.to_numpy(association_sim["index"])
0136     particle_indices = ak.to_numpy(sim_particle["index"])
0137     weights = ak.to_numpy(weights)
0138     statuses = ak.to_numpy(mc_particles["generatorStatus"])
0139     begin = ak.to_numpy(mc_particles["parents_begin"])
0140     end = ak.to_numpy(mc_particles["parents_end"])
0141     parents = ak.to_numpy(mc_parents["index"])
0142     if not (len(raw_indices) == len(sim_indices) == len(weights)):
0143         raise ValueError("RawHitAssociation weights and relation lengths disagree")
0144     if not (len(statuses) == len(begin) == len(end)):
0145         raise ValueError("MCParticle status and parent-range lengths disagree")
0146
0147     raw_to_rows = defaultdict(list)
0148     for row, raw_index in enumerate(raw_indices):
0149         if raw_index >= 0:
0150             raw_to_rows[int(raw_index)].append(row)
0151     # Cache ancestry results for particles shared by many associations/hits.
0152     resolved = {}
0153
0154     def association_status(row):
0155         sim_index = int(sim_indices[row])
0156         if sim_index < 0 or sim_index >= len(particle_indices):
0157             return None, False, False
0158         particle_index = int(particle_indices[sim_index])
0159         if particle_index not in resolved:
0160             resolved[particle_index] = resolve_status(particle_index, statuses, begin, end, parents)
0161         return resolved[particle_index]
0162
0163     labels = []
0164     diagnostics = Counter()
0165     for raw_index in ak.to_numpy(rec_raw["index"]):
0166         rows = raw_to_rows.get(int(raw_index), [])
0167         if raw_index < 0:
0168             labels.append("invalid_relation")
0169             diagnostics["invalid_relation_count"] += 1
0170             continue
0171         if not rows:
0172             labels.append("unmatched")
0173             diagnostics["unmatched_count"] += 1
0174             continue
0175         if len(rows) > 1:
0176             diagnostics["ambiguous_association_count"] += 1
0177             candidate_sources = {
0178                 STATUS_TO_SOURCE.get(status, "other")
0179                 for status, _, _ in (association_status(row) for row in rows)
0180                 if status is not None
0181             }
0182             diagnostics["ambiguous_source_count"] += int(len(candidate_sources) > 1)
0183         if not np.all(np.isfinite(weights[rows])) or np.any(weights[rows] < 0):
0184             labels.append("invalid_relation")
0185             diagnostics["invalid_weight_count"] += 1
0186             diagnostics["invalid_relation_count"] += 1
0187             continue
0188         best = max(rows, key=lambda row: float(weights[row]))
0189         status, traced, unresolved = association_status(best)
0190         if status is None:
0191             labels.append("invalid_relation")
0192             diagnostics["invalid_relation_count"] += 1
0193             continue
0194         if traced:
0195             diagnostics["secondary_unresolved_count" if unresolved else "secondary_traced_count"] += 1
0196         labels.append(STATUS_TO_SOURCE.get(status, "other"))
0197     diagnostics["rec_hits"] += len(labels)
0198     return np.asarray(labels, dtype=str), diagnostics