Examples/python/uproot.py

0001 # This file is part of the ACTS project.
0002 #
0003 # Copyright (C) 2016 CERN for the benefit of the ACTS project
0004 #
0005 # This Source Code Form is subject to the terms of the Mozilla Public
0006 # License, v. 2.0. If a copy of the MPL was not distributed with this
0007 # file, You can obtain one at https://mozilla.org/MPL/2.0/.
0008
0009 """Uproot-based readers for simulated particles and hits.
0010
0011 Reads the ROOT files written by RootParticleWriter and RootSimHitWriter
0012 using uproot, without requiring the ROOT plugin. Suitable for use in the
0013 PyPI distribution.
0014 """
0015
0016 from pathlib import Path
0017 from typing import Union
0018
0019 import numpy as np
0020 import uproot
0021
0022 import acts
0023 import acts.examples
0024
0025
0026 class UprootParticleReader(acts.examples.IReader):
0027     """Reads simulated particles from a ROOT file using uproot.
0028
0029     Reads the file format produced by RootParticleWriter (one entry per event,
0030     vector-valued branches). All data is loaded upfront for thread-safe
0031     concurrent access from the sequencer.
0032     """
0033
0034     def __init__(
0035         self,
0036         filePath: Union[Path, str],
0037         outputParticles: str = "particles",
0038         level: acts.logging.Level = acts.logging.INFO,
0039     ):
0040         acts.examples.IReader.__init__(self, "UprootParticleReader", level)
0041
0042         self._outputParticles = outputParticles
0043
0044         self._particleHandle = acts.examples.WriteDataHandle(
0045             self, acts.examples.SimParticleContainer, "OutputParticles"
0046         )
0047         self._particleHandle.initialize(self._outputParticles)
0048
0049         with uproot.open(str(filePath)) as f:
0050             tree = f["particles"]
0051             self._data = tree.arrays(library="np")
0052             event_ids = self._data["event_id"]
0053             self._entry_map = {int(eid): i for i, eid in enumerate(event_ids)}
0054
0055         self._min_event = min(self._entry_map.keys())
0056         self._max_event = max(self._entry_map.keys())
0057
0058     def availableEvents(self):
0059         return (self._min_event, self._max_event + 1)
0060
0061     def read(self, context):
0062         event_number = context.eventNumber
0063         particles = acts.examples.SimParticleContainer()
0064
0065         if event_number in self._entry_map:
0066             idx = self._entry_map[event_number]
0067             d = self._data
0068             n = len(d["particle_type"][idx])
0069             u = acts.UnitConstants
0070             for i in range(n):
0071                 barcode = acts.examples.SimBarcode()
0072                 barcode.vertexPrimary = int(d["vertex_primary"][idx][i])
0073                 barcode.vertexSecondary = int(d["vertex_secondary"][idx][i])
0074                 barcode.particle = int(d["particle"][idx][i])
0075                 barcode.generation = int(d["generation"][idx][i])
0076                 barcode.subParticle = int(d["sub_particle"][idx][i])
0077
0078                 p = acts.examples.SimParticle(
0079                     barcode,
0080                     acts.PdgParticle(int(d["particle_type"][idx][i])),
0081                     float(d["q"][idx][i]) * u.e,
0082                     float(d["m"][idx][i]) * u.GeV,
0083                 )
0084                 p.process = acts.examples.GenerationProcess(int(d["process"][idx][i]))
0085
0086                 p.fourPosition = acts.Vector4(
0087                     *[float(d[k][idx][i]) * u.mm for k in ("vx", "vy", "vz", "vt")]
0088                 )
0089                 p.direction = acts.Vector3(
0090                     *[float(d[k][idx][i]) for k in ("px", "py", "pz")]
0091                 )
0092                 p.absoluteMomentum = float(d["p"][idx][i]) * u.GeV
0093
0094                 p.setFinalMaterialPassed(
0095                     float(d["total_x0"][idx][i]) * u.mm,
0096                     float(d["total_l0"][idx][i]) * u.mm,
0097                 )
0098                 p.numberOfHits = int(d["number_of_hits"][idx][i])
0099                 p.outcome = acts.examples.SimulationOutcome(int(d["outcome"][idx][i]))
0100
0101                 particles.insert(p)
0102
0103         self._particleHandle(context, particles)
0104         return acts.examples.ProcessCode.SUCCESS
0105
0106
0107 class UprootSimHitReader(acts.examples.IReader):
0108     """Reads simulated hits from a ROOT file using uproot.
0109
0110     Reads the file format produced by RootSimHitWriter (one row per hit,
0111     scalar branches grouped by event_id). All data is loaded upfront for
0112     thread-safe concurrent access from the sequencer.
0113     """
0114
0115     def __init__(
0116         self,
0117         filePath: Union[Path, str],
0118         outputSimHits: str = "simhits",
0119         level: acts.logging.Level = acts.logging.INFO,
0120     ):
0121         acts.examples.IReader.__init__(self, "UprootSimHitReader", level)
0122
0123         self._outputSimHits = outputSimHits
0124
0125         self._hitHandle = acts.examples.WriteDataHandle(
0126             self, acts.examples.SimHitContainer, "OutputSimHits"
0127         )
0128         self._hitHandle.initialize(self._outputSimHits)
0129
0130         with uproot.open(str(filePath)) as f:
0131             tree = f["hits"]
0132             self._data = tree.arrays(library="np")
0133             self._event_range_map = self._build_event_range_map(self._data["event_id"])
0134
0135         all_ids = set(self._event_range_map.keys())
0136         self._min_event = int(min(all_ids))
0137         self._max_event = int(max(all_ids))
0138
0139     @staticmethod
0140     def _build_event_range_map(event_ids: np.ndarray) -> dict:
0141         """Build a {event_id: (start, end)} map from a sorted event_id array."""
0142         if len(event_ids) == 0:
0143             return {}
0144         unique_ids, starts = np.unique(event_ids, return_index=True)
0145         ends = np.append(starts[1:], len(event_ids))
0146         return {
0147             int(uid): (int(s), int(e)) for uid, s, e in zip(unique_ids, starts, ends)
0148         }
0149
0150     def availableEvents(self):
0151         return (self._min_event, self._max_event + 1)
0152
0153     def read(self, context):
0154         event_number = context.eventNumber
0155         hits = acts.examples.SimHitContainer()
0156
0157         if event_number in self._event_range_map:
0158             start, end = self._event_range_map[event_number]
0159             d = self._data
0160             u = acts.UnitConstants
0161             for i in range(start, end):
0162                 geoid = acts.GeometryIdentifier(int(d["geometry_id"][i]))
0163
0164                 barcode = acts.examples.SimBarcode()
0165                 barcode.vertexPrimary = int(d["barcode_vertex_primary"][i])
0166                 barcode.vertexSecondary = int(d["barcode_vertex_secondary"][i])
0167                 barcode.particle = int(d["barcode_particle"][i])
0168                 barcode.generation = int(d["barcode_generation"][i])
0169                 barcode.subParticle = int(d["barcode_sub_particle"][i])
0170
0171                 pos4 = acts.Vector4(
0172                     *[float(d[k][i]) * u.mm for k in ("tx", "ty", "tz", "tt")]
0173                 )
0174                 before4 = acts.Vector4(
0175                     *[float(d[k][i]) * u.GeV for k in ("tpx", "tpy", "tpz", "te")]
0176                 )
0177                 after4 = acts.Vector4(
0178                     (float(d["tpx"][i]) + float(d["deltapx"][i])) * u.GeV,
0179                     (float(d["tpy"][i]) + float(d["deltapy"][i])) * u.GeV,
0180                     (float(d["tpz"][i]) + float(d["deltapz"][i])) * u.GeV,
0181                     (float(d["te"][i]) + float(d["deltae"][i])) * u.GeV,
0182                 )
0183
0184                 hit = acts.examples.SimHit(
0185                     geoid, barcode, pos4, before4, after4, int(d["index"][i])
0186                 )
0187                 hits.insert(hit)
0188
0189         self._hitHandle(context, hits)
0190         return acts.examples.ProcessCode.SUCCESS