diff --git a/Python/Examples/CMakeLists.txt b/Python/Examples/CMakeLists.txt index c3ee2b8a2a1..24182d2bb9a 100644 --- a/Python/Examples/CMakeLists.txt +++ b/Python/Examples/CMakeLists.txt @@ -9,6 +9,7 @@ set(py_files reconstruction.py itk.py odd.py + uproot.py ) # Macro to add example bindings conditionally diff --git a/Python/Examples/python/uproot.py b/Python/Examples/python/uproot.py new file mode 100644 index 00000000000..b90c5df3a92 --- /dev/null +++ b/Python/Examples/python/uproot.py @@ -0,0 +1,196 @@ +# This file is part of the ACTS project. +# +# Copyright (C) 2016 CERN for the benefit of the ACTS project +# +# This Source Code Form is subject to the terms of the Mozilla Public +# License, v. 2.0. If a copy of the MPL was not distributed with this +# file, You can obtain one at https://mozilla.org/MPL/2.0/. + +"""Uproot-based readers for simulated particles and hits. + +Reads the ROOT files written by RootParticleWriter and RootSimHitWriter +using uproot, without requiring the ROOT plugin. Suitable for use in the +PyPI distribution. +""" + +from pathlib import Path +from typing import Union + +import numpy as np +import uproot + +import acts +import acts.examples + + +class UprootParticleReader(acts.examples.IReader): + """Reads simulated particles from a ROOT file using uproot. + + Reads the file format produced by RootParticleWriter (one entry per event, + vector-valued branches). All data is loaded upfront for thread-safe + concurrent access from the sequencer. + """ + + def __init__( + self, + filePath: Union[Path, str], + outputParticles: str = "particles", + level: acts.logging.Level = acts.logging.INFO, + ): + acts.examples.IReader.__init__(self, "UprootParticleReader", level) + + self._outputParticles = outputParticles + + self._particleHandle = acts.examples.WriteDataHandle( + self, acts.examples.SimParticleContainer, "OutputParticles" + ) + + with uproot.open(str(filePath)) as f: + tree = f["particles"] + self._data = tree.arrays(library="np") + event_ids = self._data["event_id"] + self._entry_map = {int(eid): i for i, eid in enumerate(event_ids)} + + self._min_event = min(self._entry_map.keys()) + self._max_event = max(self._entry_map.keys()) + + def availableEvents(self): + return (self._min_event, self._max_event + 1) + + def initialize(self): + self._particleHandle.initialize(self._outputParticles) + return acts.examples.ProcessCode.SUCCESS + + def read(self, context): + event_number = context.eventNumber + particles = acts.examples.SimParticleContainer() + + if event_number in self._entry_map: + idx = self._entry_map[event_number] + d = self._data + n = len(d["particle_type"][idx]) + u = acts.UnitConstants + for i in range(n): + barcode = acts.examples.SimBarcode() + barcode.vertexPrimary = int(d["vertex_primary"][idx][i]) + barcode.vertexSecondary = int(d["vertex_secondary"][idx][i]) + barcode.particle = int(d["particle"][idx][i]) + barcode.generation = int(d["generation"][idx][i]) + barcode.subParticle = int(d["sub_particle"][idx][i]) + + p = acts.examples.SimParticle( + barcode, + acts.PdgParticle(int(d["particle_type"][idx][i])), + float(d["q"][idx][i]) * u.e, + float(d["m"][idx][i]) * u.GeV, + ) + p.process = acts.examples.GenerationProcess(int(d["process"][idx][i])) + + p.fourPosition = acts.Vector4( + *[float(d[k][idx][i]) * u.mm for k in ("vx", "vy", "vz", "vt")] + ) + p.direction = acts.Vector3( + *[float(d[k][idx][i]) for k in ("px", "py", "pz")] + ) + p.absoluteMomentum = float(d["p"][idx][i]) * u.GeV + + p.setFinalMaterialPassed( + float(d["total_x0"][idx][i]) * u.mm, + float(d["total_l0"][idx][i]) * u.mm, + ) + p.numberOfHits = int(d["number_of_hits"][idx][i]) + p.outcome = acts.examples.SimulationOutcome(int(d["outcome"][idx][i])) + + particles.insert(p) + + self._particleHandle(context, particles) + return acts.examples.ProcessCode.SUCCESS + + +class UprootSimHitReader(acts.examples.IReader): + """Reads simulated hits from a ROOT file using uproot. + + Reads the file format produced by RootSimHitWriter (one row per hit, + scalar branches grouped by event_id). All data is loaded upfront for + thread-safe concurrent access from the sequencer. + """ + + def __init__( + self, + filePath: Union[Path, str], + outputSimHits: str = "simhits", + level: acts.logging.Level = acts.logging.INFO, + ): + acts.examples.IReader.__init__(self, "UprootSimHitReader", level) + + self._outputSimHits = outputSimHits + + self._hitHandle = acts.examples.WriteDataHandle( + self, acts.examples.SimHitContainer, "OutputSimHits" + ) + + with uproot.open(str(filePath)) as f: + tree = f["hits"] + self._data = tree.arrays(library="np") + self._event_range_map = self._build_event_range_map(self._data["event_id"]) + + all_ids = set(self._event_range_map.keys()) + self._min_event = int(min(all_ids)) + self._max_event = int(max(all_ids)) + + @staticmethod + def _build_event_range_map(event_ids: np.ndarray) -> dict: + """Build a {event_id: (start, end)} map from a sorted event_id array.""" + if len(event_ids) == 0: + return {} + unique_ids, starts = np.unique(event_ids, return_index=True) + ends = np.append(starts[1:], len(event_ids)) + return { + int(uid): (int(s), int(e)) for uid, s, e in zip(unique_ids, starts, ends) + } + + def availableEvents(self): + return (self._min_event, self._max_event + 1) + + def initialize(self): + self._hitHandle.initialize(self._outputSimHits) + return acts.examples.ProcessCode.SUCCESS + + def read(self, context): + event_number = context.eventNumber + hits = acts.examples.SimHitContainer() + + if event_number in self._event_range_map: + start, end = self._event_range_map[event_number] + d = self._data + u = acts.UnitConstants + for i in range(start, end): + geoid = acts.GeometryIdentifier(int(d["geometry_id"][i])) + + barcode = acts.examples.SimBarcode() + barcode.vertexPrimary = int(d["barcode_vertex_primary"][i]) + barcode.vertexSecondary = int(d["barcode_vertex_secondary"][i]) + barcode.particle = int(d["barcode_particle"][i]) + barcode.generation = int(d["barcode_generation"][i]) + barcode.subParticle = int(d["barcode_sub_particle"][i]) + + pos4 = acts.Vector4( + *[float(d[k][i]) * u.mm for k in ("tx", "ty", "tz", "tt")] + ) + before4 = acts.Vector4( + *[float(d[k][i]) * u.GeV for k in ("tpx", "tpy", "tpz", "te")] + ) + after4 = acts.Vector4( + (float(d["tpx"][i]) + float(d["deltapx"][i])) * u.GeV, + (float(d["tpy"][i]) + float(d["deltapy"][i])) * u.GeV, + (float(d["tpz"][i]) + float(d["deltapz"][i])) * u.GeV, + (float(d["te"][i]) + float(d["deltae"][i])) * u.GeV, + ) + + hit = acts.examples.SimHit( + geoid, barcode, pos4, before4, after4, int(d["index"][i]) + ) + hits.insert(hit) + + self._hitHandle(context, hits) + return acts.examples.ProcessCode.SUCCESS diff --git a/Python/Examples/src/Framework.cpp b/Python/Examples/src/Framework.cpp index 8b233e97cef..7088aafed62 100644 --- a/Python/Examples/src/Framework.cpp +++ b/Python/Examples/src/Framework.cpp @@ -101,6 +101,47 @@ class PyIAlgorithm : public IAlgorithm { const Acts::Logger& pyLogger() const { return logger(); } }; +class PyIReader : public IReader { + // Typedef avoids comma in macro argument for PYBIND11_OVERRIDE_PURE + using EventRange = std::pair; + + public: + PyIReader(std::string name, Logging::Level level) + : IReader(), + m_name(std::move(name)), + m_logger(Acts::getDefaultLogger(m_name, level)) {} + + std::string name() const override { return m_name; } + + std::string_view typeName() const override { return "Reader"; } + + EventRange availableEvents() const override { + py::gil_scoped_acquire acquire{}; + PYBIND11_OVERRIDE_PURE(EventRange, IReader, availableEvents); + } + + ProcessCode read(const AlgorithmContext& ctx) override { + py::gil_scoped_acquire acquire{}; + PYBIND11_OVERRIDE_PURE(ProcessCode, IReader, read, ctx); + } + + ProcessCode initialize() override { + py::gil_scoped_acquire acquire{}; + PYBIND11_OVERRIDE(ProcessCode, IReader, initialize); + } + + ProcessCode finalize() override { + py::gil_scoped_acquire acquire{}; + PYBIND11_OVERRIDE(ProcessCode, IReader, finalize); + } + + const Acts::Logger& pyLogger() const { return *m_logger; } + + private: + std::string m_name; + std::unique_ptr m_logger; +}; + class PyReadDataHandle : public ReadDataHandleBase { public: PyReadDataHandle(SequenceElement* parent, py::object pytype, @@ -206,8 +247,6 @@ namespace ActsPython { void addFramework(py::module& mex) { py::class_>(mex, "IWriter"); - py::class_>(mex, "IReader"); - py::class_>( mex, "IContextDecorator") .def("decorate", &IContextDecorator::decorate) @@ -283,6 +322,18 @@ void addFramework(py::module& mex) { .def("internalExecute", &SequenceElement::internalExecute) .def("name", &SequenceElement::name); + py::class_, SequenceElement, PyIReader>( + mex, "IReader") + .def(py::init_alias(), py::arg("name"), + py::arg("level")) + .def("availableEvents", &IReader::availableEvents) + .def("read", &IReader::read) + .def("name", &IReader::name) + .def_property_readonly("logger", + [](const PyIReader& self) -> const Acts::Logger& { + return self.pyLogger(); + }); + auto bareAlgorithm = py::class_, SequenceElement, PyIAlgorithm>(mex, "IAlgorithm") diff --git a/Python/Examples/src/Generators.cpp b/Python/Examples/src/Generators.cpp index 3fc8c65ae8a..b60389985ac 100644 --- a/Python/Examples/src/Generators.cpp +++ b/Python/Examples/src/Generators.cpp @@ -8,16 +8,20 @@ #include "Acts/Definitions/Algebra.hpp" #include "Acts/Definitions/PdgParticle.hpp" +#include "Acts/Geometry/GeometryIdentifier.hpp" #include "Acts/Utilities/AngleHelpers.hpp" #include "Acts/Utilities/Logger.hpp" +#include "ActsExamples/EventData/SimHit.hpp" #include "ActsExamples/EventData/SimParticle.hpp" #include "ActsExamples/Generators/EventGenerator.hpp" #include "ActsExamples/Utilities/ParametricParticleGenerator.hpp" #include "ActsExamples/Utilities/VertexGenerators.hpp" +#include "ActsFatras/EventData/Hit.hpp" #include "ActsPython/Utilities/Macros.hpp" #include "ActsPython/Utilities/WhiteBoardRegistry.hpp" #include +#include #include #include #include @@ -155,33 +159,63 @@ void addGenerators(py::module& mex) { py::arg("pdg")) .def(py::init(), py::arg("initial"), py::arg("final")) - .def_property_readonly("particleId", &SimParticle::particleId) - .def_property_readonly("pdg", &SimParticle::pdg) + .def_property("particleId", &SimParticle::particleId, + [](SimParticle& p, SimBarcode b) { p.setParticleId(b); }) + .def_property("pdg", &SimParticle::pdg, + [](SimParticle& p, Acts::PdgParticle v) { p.setPdg(v); }) .def_property_readonly("absolutePdg", &SimParticle::absolutePdg) - .def_property_readonly("charge", &SimParticle::charge) + .def_property("charge", &SimParticle::charge, + [](SimParticle& p, double v) { p.setCharge(v); }) .def_property_readonly("absoluteCharge", &SimParticle::absoluteCharge) - .def_property_readonly("mass", &SimParticle::mass) - .def_property_readonly("process", &SimParticle::process) + .def_property("mass", &SimParticle::mass, + [](SimParticle& p, double v) { p.setMass(v); }) + .def_property("process", &SimParticle::process, + [](SimParticle& p, ActsFatras::GenerationProcess v) { + p.setProcess(v); + }) .def_property_readonly("isSecondary", &SimParticle::isSecondary) - .def_property_readonly("fourPosition", &SimParticle::fourPosition) + .def_property("fourPosition", &SimParticle::fourPosition, + [](SimParticle& p, const Acts::Vector4& pos4) { + p.initialState().setPosition4(pos4); + }) .def_property_readonly( "position", [](const SimParticle& p) { return Vector3(p.position()); }) .def_property_readonly("time", &SimParticle::time) .def_property_readonly("fourMomentum", &SimParticle::fourMomentum) - .def_property_readonly("direction", &SimParticle::direction) + .def_property("direction", &SimParticle::direction, + [](SimParticle& p, const Acts::Vector3& dir) { + p.initialState().setDirection(dir); + }) .def_property_readonly("theta", &SimParticle::theta) .def_property_readonly("phi", &SimParticle::phi) .def_property_readonly("transverseMomentum", &SimParticle::transverseMomentum) - .def_property_readonly("absoluteMomentum", &SimParticle::absoluteMomentum) + .def_property("absoluteMomentum", &SimParticle::absoluteMomentum, + [](SimParticle& p, double v) { + p.initialState().setAbsoluteMomentum(v); + }) .def_property_readonly("momentum", &SimParticle::momentum) .def_property_readonly("energy", &SimParticle::energy) .def_property_readonly("energyLoss", &SimParticle::energyLoss) .def_property_readonly("pathInX0", &SimParticle::pathInX0) .def_property_readonly("pathInL0", &SimParticle::pathInL0) - .def_property_readonly("numberOfHits", &SimParticle::numberOfHits) - .def_property_readonly("outcome", &SimParticle::outcome) + .def( + "setFinalMaterialPassed", + [](SimParticle& p, double x0, double l0) -> SimParticle& { + p.finalState().setMaterialPassed(x0, l0); + return p; + }, + py::arg("pathInX0"), py::arg("pathInL0"), + py::return_value_policy::reference_internal) + .def_property("numberOfHits", &SimParticle::numberOfHits, + [](SimParticle& p, std::uint32_t n) { + p.finalState().setNumberOfHits(n); + }) + .def_property("outcome", &SimParticle::outcome, + [](SimParticle& p, ActsFatras::SimulationOutcome o) { + p.finalState().setOutcome(o); + }) .def_property_readonly( "initialState", py::overload_cast<>(&SimParticle::initialState, py::const_)) @@ -211,13 +245,56 @@ void addGenerators(py::module& mex) { [](const SimParticleContainer& c, SimBarcode barcode) { return c.find(barcode) != c.end(); }) - .def("__repr__", [](const SimParticleContainer& c) { + .def("__repr__", + [](const SimParticleContainer& c) { + std::ostringstream oss; + oss << "SimParticleContainer(" << c.size() << " particles)"; + return oss.str(); + }) + .def( + "insert", + [](SimParticleContainer& c, const SimParticle& p) { + c.insert(p); + }, + py::arg("particle")); + + WhiteBoardRegistry::registerClass(simParticleContainer); + + py::class_(mex, "SimHit") + .def(py::init(), + py::arg("geometryId"), py::arg("particleId"), py::arg("pos4"), + py::arg("before4"), py::arg("after4"), py::arg("index") = -1) + .def_property_readonly("geometryId", &SimHit::geometryId) + .def_property_readonly("particleId", &SimHit::particleId) + .def_property_readonly("index", &SimHit::index) + .def_property_readonly("fourPosition", &SimHit::fourPosition) + .def_property_readonly("time", &SimHit::time) + .def_property_readonly("momentum4Before", &SimHit::momentum4Before) + .def_property_readonly("momentum4After", &SimHit::momentum4After) + .def_property_readonly("depositedEnergy", &SimHit::depositedEnergy); + + auto simHitContainer = + py::classh(mex, "SimHitContainer") + .def(py::init<>()) + .def("__len__", [](const SimHitContainer& c) { return c.size(); }) + .def( + "__iter__", + [](const SimHitContainer& c) { + return py::make_iterator(c.begin(), c.end()); + }, + py::keep_alive<0, 1>()) + .def( + "insert", + [](SimHitContainer& c, const SimHit& h) { c.insert(h); }, + py::arg("hit")) + .def("__repr__", [](const SimHitContainer& c) { std::ostringstream oss; - oss << "SimParticleContainer(" << c.size() << " particles)"; + oss << "SimHitContainer(" << c.size() << " hits)"; return oss.str(); }); - WhiteBoardRegistry::registerClass(simParticleContainer); + WhiteBoardRegistry::registerClass(simHitContainer); { using Config = ParametricParticleGenerator::Config; diff --git a/Python/Examples/tests/test_uproot.py b/Python/Examples/tests/test_uproot.py new file mode 100644 index 00000000000..f93b37ed1cf --- /dev/null +++ b/Python/Examples/tests/test_uproot.py @@ -0,0 +1,69 @@ +# This file is part of the ACTS project. +# +# Copyright (C) 2016 CERN for the benefit of the ACTS project +# +# This Source Code Form is subject to the terms of the Mozilla Public +# License, v. 2.0. If a copy of the MPL was not distributed with this +# file, You can obtain one at https://mozilla.org/MPL/2.0/. + +import pytest +from pathlib import Path + +from helpers import AssertCollectionExistsAlg + +import acts +import acts.examples +from acts.examples.uproot import UprootParticleReader, UprootSimHitReader +from acts.examples import Sequencer +from acts.examples.root import RootParticleWriter, RootSimHitWriter + + +@pytest.mark.root +def test_root_write_uproot_read(tmp_path, fatras, conf_const): + """Full round-trip: simulate with fatras, write ROOT files, read back with uproot readers.""" + particle_file = tmp_path / "particles_simulation.root" + hit_file = tmp_path / "hits.root" + + s = Sequencer(numThreads=1, events=10, logLevel=acts.logging.WARNING) + evGen, simAlg, _ = fatras(s) + s.addWriter( + conf_const( + RootParticleWriter, + acts.logging.WARNING, + inputParticles=simAlg.config.outputParticles, + filePath=str(particle_file), + ) + ) + s.addWriter( + conf_const( + RootSimHitWriter, + acts.logging.WARNING, + inputSimHits=simAlg.config.outputSimHits, + filePath=str(hit_file), + ) + ) + s.run() + + s2 = Sequencer(numThreads=1, logLevel=acts.logging.WARNING) + s2.addReader( + UprootParticleReader( + filePath=particle_file, + outputParticles="particles_read", + level=acts.logging.WARNING, + ) + ) + s2.addReader( + UprootSimHitReader( + filePath=hit_file, + outputSimHits="simhits_read", + level=acts.logging.WARNING, + ) + ) + alg = AssertCollectionExistsAlg( + ["particles_read", "simhits_read"], + "check_alg", + acts.logging.WARNING, + ) + s2.addAlgorithm(alg) + s2.run() + assert alg.events_seen == 10