pyrobird/pyrobird/edm4hep.py

0001 import logging
0002 import math
0003 from collections import defaultdict
0004
0005 import awkward as ak
0006
0007 from pyrobird.dex import make_dex, PIECE_VERSION
0008 from pyrobird.mc_particles import (mc_particles_to_trajectories,
0009                                    DEFAULT_MC_STEP_TIME, DEFAULT_MC_MAX_POINTS)
0010
0011 """
0012 Converts EDM4hep simulation output (ddsim / DD4hep) to Firebird DEX format.
0013
0014 Supported collections:
0015     vector<edm4hep::SimTrackerHitData> - simulated tracker hits, converted to BoxHit groups
0016     vector<edm4hep::MCParticleData>    - MC particles, used to build MC-truth trajectories
0017                                          connecting the sim hits of each particle
0018                                          ('mc_trajectories'), and straight vertex->endpoint
0019                                          lines for every particle ('mc_particles',
0020                                          see pyrobird.mc_particles)
0021
0022 SimTrackerHitData branch layout (edm4hep >= 0.99; older versions use 'EDep' instead of 'eDep'
0023 and name the particle relation '_<Collection>_MCParticle' instead of '_<Collection>_particle'):
0024     'SiBarrelHits': 'vector<edm4hep::SimTrackerHitData>',
0025     'SiBarrelHits/SiBarrelHits.cellID': 'uint64_t[]',
0026     'SiBarrelHits/SiBarrelHits.eDep': 'float[]',
0027     'SiBarrelHits/SiBarrelHits.time': 'float[]',
0028     'SiBarrelHits/SiBarrelHits.pathLength': 'float[]',
0029     'SiBarrelHits/SiBarrelHits.quality': 'int32_t[]',
0030     'SiBarrelHits/SiBarrelHits.position.x': 'double[]',
0031     'SiBarrelHits/SiBarrelHits.position.y': 'double[]',
0032     'SiBarrelHits/SiBarrelHits.position.z': 'double[]',
0033     'SiBarrelHits/SiBarrelHits.momentum.x': 'float[]',
0034     'SiBarrelHits/SiBarrelHits.momentum.y': 'float[]',
0035     'SiBarrelHits/SiBarrelHits.momentum.z': 'float[]',
0036     '_SiBarrelHits_particle': 'vector<podio::ObjectID>',
0037     '_SiBarrelHits_particle/_SiBarrelHits_particle.index': 'int32_t[]',
0038     '_SiBarrelHits_particle/_SiBarrelHits_particle.collectionID': 'uint32_t[]',
0039
0040 Sim hits carry no positionError/timeError/edepError, so BoxHit dimensions use a fixed,
0041 user-configurable box size and errors are written as 0.
0042 """
0043
0044 logger = logging.getLogger(__name__)
0045
0046 # Sim hits have no positionError, so boxes get a fixed visualization size [mm]
0047 DEFAULT_HIT_BOX_SIZE = 2.0
0048
0049 # Cherenkov/PID collections record photon detections attributed to the emitting charged
0050 # particle (optical photons are not stored in MCParticles). Connecting them into that
0051 # particle's trajectory draws zigzags along the photosensor planes, so they are excluded
0052 # from trajectory building by default (they are still converted as BoxHit groups).
0053 DEFAULT_TRAJECTORY_EXCLUDED_COLLECTIONS = ("DIRCBarHits", "DRICHHits", "PFRICHHits")
0054
0055
0056 def get_sim_hit_branches(tree):
0057     """Returns names of all vector<edm4hep::SimTrackerHitData> branches in the tree"""
0058     typenames = tree.typenames(recursive=False, full_paths=True,
0059                                filter_typename="vector<edm4hep::SimTrackerHitData>")
0060     return list(typenames.keys())
0061
0062
0063 def detect_file_type(tree):
0064     """
0065     Detects the data model of an opened 'events' tree by its branch types.
0066
0067     Returns 'edm4eic' if reconstructed edm4eic::TrackerHitData collections are present
0068     (eicrecon files also contain sim hits, reconstruction wins), 'edm4hep' if only
0069     edm4hep::SimTrackerHitData collections are present.
0070
0071     Raises
0072     ------
0073     ValueError
0074         If neither collection type is found in the tree.
0075     """
0076     if tree.typenames(recursive=False, full_paths=True, filter_typename="vector<edm4eic::TrackerHitData>"):
0077         return "edm4eic"
0078     if tree.typenames(recursive=False, full_paths=True, filter_typename="vector<edm4hep::SimTrackerHitData>"):
0079         return "edm4hep"
0080     raise ValueError("Cannot detect file type: no 'vector<edm4eic::TrackerHitData>' or "
0081                      "'vector<edm4hep::SimTrackerHitData>' branches found in 'events' tree")
0082
0083
0084 def _get_field_array(tree, full_branch, entry_start, entry_stop):
0085     """Reads a jagged field branch and returns it flattened as a python list"""
0086     return ak.flatten(tree[full_branch].array(entry_start=entry_start, entry_stop=entry_stop)).to_list()
0087
0088
0089 def _find_particle_relation_branch(tree, branch_name):
0090     """Finds the hit->MCParticle relation branch name, or None if missing.
0091
0092     edm4hep >= 0.99 names it '_<Collection>_particle', older versions '_<Collection>_MCParticle'.
0093     """
0094     keys = set(tree.keys(recursive=False))
0095     for suffix in ("particle", "MCParticle"):
0096         candidate = f"_{branch_name}_{suffix}"
0097         if candidate in keys:
0098             return candidate
0099     return None
0100
0101
0102 def _find_edep_field(tree, branch_name):
0103     """Returns the energy deposit field name: 'eDep' (edm4hep >= 0.99) or 'EDep' (older)"""
0104     keys = set(tree[branch_name].keys())
0105     for field in ("eDep", "EDep"):
0106         if f"{branch_name}.{field}" in keys:
0107             return field
0108     return None
0109
0110
0111 def _get_mc_collection_id(tree, mc_branch):
0112     """
0113     Looks up the podio collectionID of `mc_branch` from the file's podio_metadata tree.
0114     Returns None if the metadata is unavailable (e.g. trimmed test files).
0115     """
0116     try:
0117         metadata_tree = tree.file["podio_metadata"]
0118         info_branch = "events___CollectionTypeInfo"
0119         ids = metadata_tree[f"{info_branch}/{info_branch}.collectionID"].array().tolist()[0]
0120         names = metadata_tree[f"{info_branch}/{info_branch}.name"].array().tolist()[0]
0121         for coll_id, name in zip(ids, names):
0122             if name == mc_branch:
0123                 return coll_id
0124     except Exception as ex:
0125         logger.debug(f"podio_metadata not available for collectionID check: {ex}")
0126     return None
0127
0128
0129 def sim_tracker_hits_to_box_hits(tree, branch_name, entry_start, entry_stop=None,
0130                                  box_size=DEFAULT_HIT_BOX_SIZE):
0131     """Converts vector<edm4hep::SimTrackerHitData> to BoxHit format dictionary.
0132
0133     Sim hits have no positionError, so all boxes use a fixed `box_size` [mm] cube
0134     and time/energy errors are written as 0.
0135     """
0136
0137     # Read only 1 event if entry_stop is not given
0138     if entry_stop is None:
0139         entry_stop = entry_start + 1
0140
0141     def get_field_array(field):
0142         return _get_field_array(tree, f"{branch_name}/{branch_name}.{field}", entry_start, entry_stop)
0143
0144     pos_x = get_field_array("position.x")   # 'double[]'
0145     pos_y = get_field_array("position.y")   # 'double[]'
0146     pos_z = get_field_array("position.z")   # 'double[]'
0147     time = get_field_array("time")          # 'float[]'
0148
0149     edep_field = _find_edep_field(tree, branch_name)
0150     if edep_field:
0151         edep = get_field_array(edep_field)
0152     else:
0153         logger.warning(f"No eDep/EDep field found in '{branch_name}', writing 0 energy deposit")
0154         edep = [0.0] * len(pos_x)
0155
0156     # Columnar piece: parallel arrays, hit id == array index; pos is flat xyz triplets.
0157     # Sim data carries no errors, so the error columns are omitted entirely.
0158     count = len(pos_x)
0159     pos = []
0160     for i in range(count):
0161         pos.extend((pos_x[i], pos_y[i], pos_z[i]))
0162
0163     piece = {
0164         "name": branch_name,
0165         "type": "BoxHit",
0166         "version": PIECE_VERSION,
0167         "origin": {"type": "edm4hep::SimTrackerHitData", "name": branch_name},
0168         "count": count,
0169         "columns": {
0170             "pos": pos,
0171             "dim": [box_size] * (3 * count),
0172             "time": time,
0173             "edep": edep,
0174         },
0175     }
0176     return piece
0177
0178
0179 def sim_hits_to_trajectories(tree, hit_branch_names, entry_start, entry_stop=None,
0180                              mc_branch="MCParticles", min_hits=2,
0181                              prepend_vertex=True, append_endpoint=False):
0182     """
0183     Builds MC-truth trajectories connecting edm4hep::SimTrackerHitData hits.
0184
0185     Hits from all `hit_branch_names` collections are pooled, grouped by their associated
0186     MCParticle and sorted by time. Each particle with at least `min_hits` hits becomes one
0187     trajectory in a single 'PointTrajectory' group.
0188
0189     Parameters
0190     ----------
0191     tree - uproot 'events' tree
0192     hit_branch_names - SimTrackerHitData collection names to pool hits from
0193     entry_start - first entry to process
0194     entry_stop - one past last entry (defaults to entry_start + 1, i.e. a single event)
0195     mc_branch - name of the MCParticle collection the hit relations point to
0196     min_hits - particles with fewer hits are skipped (their hits stay visible as boxes)
0197     prepend_vertex - add MCParticles.vertex (with the particle creation time) as point 0
0198     append_endpoint - add MCParticles.endpoint as the last point; edm4hep stores no endpoint
0199                       time, so the point reuses the last hit's time
0200     """
0201     if entry_stop is None:
0202         entry_stop = entry_start + 1
0203
0204     # Columnar piece: track parameters are per-name columns (trajectory id == index),
0205     # the ragged per-trajectory point lists stay nested under "points"
0206     result = {
0207         "name": "McHitTrajectories",
0208         "type": "PointTrajectory",
0209         "version": PIECE_VERSION,
0210         "origin": {"type": ["edm4hep::SimTrackerHitData", "edm4hep::MCParticleData"],
0211                    "collections": list(hit_branch_names)},
0212         "count": 0,
0213         "columns": {},
0214         "pointColumns": ["x", "y", "z", "t", "dx", "dy", "dz", "dt"],
0215         "points": []
0216     }
0217
0218     if mc_branch not in tree.keys(recursive=False):
0219         logger.warning(f"'{mc_branch}' collection not found, no trajectories are built")
0220         return result
0221
0222     def get_mc_field(field):
0223         return _get_field_array(tree, f"{mc_branch}/{mc_branch}.{field}", entry_start, entry_stop)
0224
0225     mc_pdg = get_mc_field("PDG")
0226     mc_charge = get_mc_field("charge")
0227     mc_time = get_mc_field("time")
0228     mc_mom_x = get_mc_field("momentum.x")
0229     mc_mom_y = get_mc_field("momentum.y")
0230     mc_mom_z = get_mc_field("momentum.z")
0231     mc_vtx_x = get_mc_field("vertex.x")
0232     mc_vtx_y = get_mc_field("vertex.y")
0233     mc_vtx_z = get_mc_field("vertex.z")
0234     mc_end_x = get_mc_field("endpoint.x") if append_endpoint else []
0235     mc_end_y = get_mc_field("endpoint.y") if append_endpoint else []
0236     mc_end_z = get_mc_field("endpoint.z") if append_endpoint else []
0237
0238     # eicrecon files carry a second MCParticle collection (MCParticlesHeadOnFrameNoBeamFX);
0239     # all known sim-hit relations point to 'MCParticles', but sanity-check via collectionID
0240     # when podio metadata is available
0241     mc_collection_id = _get_mc_collection_id(tree, mc_branch)
0242
0243     # Pool hits from all collections grouped by MCParticle index
0244     # Each hit is a (time, x, y, z) tuple
0245     hits_by_particle = defaultdict(list)
0246     for branch_name in hit_branch_names:
0247         relation_branch = _find_particle_relation_branch(tree, branch_name)
0248         if relation_branch is None:
0249             logger.warning(f"No MCParticle relation branch found for '{branch_name}', "
0250                            f"its hits are skipped in trajectory building")
0251             continue
0252
0253         def get_field_array(field):
0254             return _get_field_array(tree, f"{branch_name}/{branch_name}.{field}", entry_start, entry_stop)
0255
0256         pos_x = get_field_array("position.x")
0257         pos_y = get_field_array("position.y")
0258         pos_z = get_field_array("position.z")
0259         time = get_field_array("time")
0260         mc_index = _get_field_array(tree, f"{relation_branch}/{relation_branch}.index", entry_start, entry_stop)
0261         coll_ids = _get_field_array(tree, f"{relation_branch}/{relation_branch}.collectionID", entry_start, entry_stop)
0262
0263         for i in range(len(pos_x)):
0264             index = mc_index[i]
0265             if index < 0 or index >= len(mc_pdg):
0266                 continue  # unfilled relation
0267             if mc_collection_id is not None and coll_ids[i] != mc_collection_id:
0268                 logger.warning(f"Hit in '{branch_name}' points to an unexpected MCParticle "
0269                                f"collection (collectionID={coll_ids[i]}), skipping it")
0270                 continue
0271             hits_by_particle[index].append((time[i], pos_x[i], pos_y[i], pos_z[i]))
0272
0273     columns = {name: [] for name in ("pdg", "charge", "px", "py", "pz", "p", "mc_index")}
0274     all_points = []
0275     for index in sorted(hits_by_particle.keys()):
0276         particle_hits = hits_by_particle[index]
0277         if len(particle_hits) < min_hits:
0278             continue
0279         particle_hits.sort(key=lambda hit: hit[0])
0280
0281         # pointColumns => [x, y, z, t, dx, dy, dz, dt]; sim data has no errors
0282         points = []
0283         if prepend_vertex:
0284             points.append([mc_vtx_x[index], mc_vtx_y[index], mc_vtx_z[index], mc_time[index], 0, 0, 0, 0])
0285         for hit_time, x, y, z in particle_hits:
0286             points.append([x, y, z, hit_time, 0, 0, 0, 0])
0287         if append_endpoint:
0288             # edm4hep stores no endpoint time, reuse the last hit time to keep time animation monotonic
0289             last_time = particle_hits[-1][0]
0290             points.append([mc_end_x[index], mc_end_y[index], mc_end_z[index], last_time, 0, 0, 0, 0])
0291
0292         momentum = math.sqrt(mc_mom_x[index] ** 2 + mc_mom_y[index] ** 2 + mc_mom_z[index] ** 2)
0293         columns["pdg"].append(mc_pdg[index])
0294         columns["charge"].append(mc_charge[index])
0295         columns["px"].append(mc_mom_x[index])
0296         columns["py"].append(mc_mom_y[index])
0297         columns["pz"].append(mc_mom_z[index])
0298         columns["p"].append(momentum)
0299         columns["mc_index"].append(index)
0300         all_points.append(points)
0301
0302     result["count"] = len(all_points)
0303     result["columns"] = columns if all_points else {}
0304     result["points"] = all_points
0305     return result
0306
0307
0308 def edm4hep_entry_to_dict(tree, entry_index, custom_name=None, collections=None,
0309                           box_size=DEFAULT_HIT_BOX_SIZE,
0310                           prepend_vertex=True, append_endpoint=False,
0311                           trajectory_excluded_collections=DEFAULT_TRAJECTORY_EXCLUDED_COLLECTIONS,
0312                           mc_step_time=DEFAULT_MC_STEP_TIME, mc_max_points=DEFAULT_MC_MAX_POINTS):
0313     # the result of this function
0314     components = []
0315
0316     if not collections:
0317         collections = [
0318             "tracker_hits",
0319             "mc_trajectories",
0320             "mc_particles",
0321         ]
0322
0323     hit_branches = get_sim_hit_branches(tree)
0324
0325     # Hits:
0326     if "tracker_hits" in collections:
0327         for branch_name in hit_branches:
0328             piece = sim_tracker_hits_to_box_hits(tree, branch_name, entry_index, box_size=box_size)
0329             # Collections empty in this event are skipped
0330             if piece["count"] > 0:
0331                 components.append(piece)
0332
0333     # MC-truth trajectories connecting the hits
0334     if "mc_trajectories" in collections:
0335         trajectory_branches = [b for b in hit_branches if b not in trajectory_excluded_collections]
0336         trajectory_piece = sim_hits_to_trajectories(tree, trajectory_branches, entry_index,
0337                                                     prepend_vertex=prepend_vertex,
0338                                                     append_endpoint=append_endpoint)
0339         if trajectory_piece["count"] > 0:
0340             components.append(trajectory_piece)
0341
0342     # Straight vertex->endpoint lines for every MCParticle
0343     if "mc_particles" in collections:
0344         mc_piece = mc_particles_to_trajectories(tree, entry_index,
0345                                                 step_time=mc_step_time, max_points=mc_max_points)
0346         if mc_piece["count"] > 0:
0347             components.append(mc_piece)
0348
0349     entry = {
0350         "id": custom_name if custom_name else entry_index,
0351         "pieces": components
0352     }
0353
0354     return entry
0355
0356
0357 def edm4hep_to_dex_dict(tree, event_ids, origin_info=None, collections=None,
0358                         box_size=DEFAULT_HIT_BOX_SIZE,
0359                         prepend_vertex=True, append_endpoint=False,
0360                         trajectory_excluded_collections=DEFAULT_TRAJECTORY_EXCLUDED_COLLECTIONS,
0361                         mc_step_time=DEFAULT_MC_STEP_TIME, mc_max_points=DEFAULT_MC_MAX_POINTS):
0362     event_data = []
0363
0364     if isinstance(event_ids, int):
0365         event_ids = [event_ids]
0366
0367     for entry_id in event_ids:
0368         event_data.append(edm4hep_entry_to_dict(
0369             tree, entry_id, custom_name=None, collections=collections,
0370             box_size=box_size, prepend_vertex=prepend_vertex, append_endpoint=append_endpoint,
0371             trajectory_excluded_collections=trajectory_excluded_collections,
0372             mc_step_time=mc_step_time, mc_max_points=mc_max_points))
0373
0374     return make_dex(event_data, origin_info)