Root/src/RootParticleWriter.cpp

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 #include "ActsExamples/Io/Root/RootParticleWriter.hpp"
0010
0011 #include "Acts/Definitions/TrackParametrization.hpp"
0012 #include "Acts/Definitions/Units.hpp"
0013 #include "Acts/Propagator/Propagator.hpp"
0014 #include "Acts/Propagator/SympyStepper.hpp"
0015 #include "Acts/Surfaces/PerigeeSurface.hpp"
0016 #include "Acts/Surfaces/Surface.hpp"
0017 #include "Acts/Utilities/Helpers.hpp"
0018 #include "Acts/Utilities/VectorHelpers.hpp"
0019 #include "ActsExamples/EventData/SimParticle.hpp"
0020 #include "ActsExamples/Framework/AlgorithmContext.hpp"
0021
0022 #include <cstdint>
0023 #include <ios>
0024 #include <stdexcept>
0025
0026 #include <TFile.h>
0027 #include <TTree.h>
0028
0029 namespace ActsExamples {
0030
0031 RootParticleWriter::RootParticleWriter(const RootParticleWriter::Config& cfg,
0032                                        Acts::Logging::Level lvl)
0033     : WriterT(cfg.inputParticles, "RootParticleWriter", lvl), m_cfg(cfg) {
0034   // inputParticles is already checked by base constructor
0035   if (m_cfg.filePath.empty()) {
0036     throw std::invalid_argument("Missing file path");
0037   }
0038   if (m_cfg.treeName.empty()) {
0039     throw std::invalid_argument("Missing tree name");
0040   }
0041
0042   // open root file and create the tree
0043   m_outputFile = TFile::Open(m_cfg.filePath.c_str(), m_cfg.fileMode.c_str());
0044   if (m_outputFile == nullptr) {
0045     throw std::ios_base::failure("Could not open '" + m_cfg.filePath + "'");
0046   }
0047   m_outputFile->cd();
0048   m_outputTree = new TTree(m_cfg.treeName.c_str(), m_cfg.treeName.c_str());
0049   if (m_outputTree == nullptr) {
0050     throw std::bad_alloc();
0051   }
0052
0053   // setup the branches
0054   m_outputTree->Branch("event_id", &m_eventId);
0055   m_outputTree->Branch("particle_hash", &m_particleHash);
0056   m_outputTree->Branch("particle_type", &m_particleType);
0057   m_outputTree->Branch("process", &m_process);
0058   m_outputTree->Branch("vx", &m_vx);
0059   m_outputTree->Branch("vy", &m_vy);
0060   m_outputTree->Branch("vz", &m_vz);
0061   m_outputTree->Branch("vt", &m_vt);
0062   m_outputTree->Branch("px", &m_px);
0063   m_outputTree->Branch("py", &m_py);
0064   m_outputTree->Branch("pz", &m_pz);
0065   m_outputTree->Branch("m", &m_m);
0066   m_outputTree->Branch("q", &m_q);
0067   m_outputTree->Branch("eta", &m_eta);
0068   m_outputTree->Branch("phi", &m_phi);
0069   m_outputTree->Branch("pt", &m_pt);
0070   m_outputTree->Branch("p", &m_p);
0071   m_outputTree->Branch("q_over_p", &m_qop);
0072   m_outputTree->Branch("theta", &m_theta);
0073   m_outputTree->Branch("vertex_primary", &m_vertexPrimary);
0074   m_outputTree->Branch("vertex_secondary", &m_vertexSecondary);
0075   m_outputTree->Branch("particle", &m_particle);
0076   m_outputTree->Branch("generation", &m_generation);
0077   m_outputTree->Branch("sub_particle", &m_subParticle);
0078   m_outputTree->Branch("orig_part_idx", &m_origParticleIdx);
0079   m_outputTree->Branch("hf_origin", &m_hfOrigin);
0080
0081   if (m_cfg.writeHelixParameters) {
0082     m_outputTree->Branch("perigee_d0", &m_perigeeD0);
0083     m_outputTree->Branch("perigee_z0", &m_perigeeZ0);
0084     m_outputTree->Branch("perigee_phi", &m_perigeePhi);
0085     m_outputTree->Branch("perigee_theta", &m_perigeeTheta);
0086     m_outputTree->Branch("perigee_q_over_p", &m_perigeeQop);
0087     m_outputTree->Branch("perigee_p", &m_perigeeP);
0088     m_outputTree->Branch("perigee_px", &m_perigeePx);
0089     m_outputTree->Branch("perigee_py", &m_perigeePy);
0090     m_outputTree->Branch("perigee_pz", &m_perigeePz);
0091     m_outputTree->Branch("perigee_eta", &m_perigeeEta);
0092     m_outputTree->Branch("perigee_pt", &m_perigeePt);
0093   }
0094
0095   m_outputTree->Branch("e_loss", &m_eLoss);
0096   m_outputTree->Branch("total_x0", &m_pathInX0);
0097   m_outputTree->Branch("total_l0", &m_pathInL0);
0098   m_outputTree->Branch("number_of_hits", &m_numberOfHits);
0099   m_outputTree->Branch("outcome", &m_outcome);
0100 }
0101
0102 RootParticleWriter::~RootParticleWriter() {
0103   if (m_outputFile != nullptr) {
0104     m_outputFile->Close();
0105   }
0106 }
0107
0108 ProcessCode RootParticleWriter::finalize() {
0109   m_outputFile->cd();
0110   m_outputTree->Write();
0111   m_outputFile->Close();
0112
0113   ACTS_INFO("Wrote particles to tree '" << m_cfg.treeName << "' in '"
0114                                         << m_cfg.filePath << "'");
0115
0116   return ProcessCode::SUCCESS;
0117 }
0118
0119 ProcessCode RootParticleWriter::writeT(const AlgorithmContext& ctx,
0120                                        const SimParticleContainer& particles) {
0121   // ensure exclusive access to tree/file while writing
0122   std::lock_guard<std::mutex> lock(m_writeMutex);
0123
0124   m_eventId = ctx.eventNumber;
0125   for (const auto& particle : particles) {
0126     m_particleHash.push_back(particle.particleId().hash());
0127     m_particleType.push_back(particle.pdg());
0128     m_origParticleIdx.push_back(particle.origParticleIdx());
0129     m_hfOrigin.push_back(
0130         static_cast<std::uint8_t>(particle.heavyFlavourOrigin()));
0131     m_process.push_back(static_cast<std::uint32_t>(particle.process()));
0132     // position
0133     m_vx.push_back(Acts::clampValue<float>(particle.fourPosition().x() /
0134                                            Acts::UnitConstants::mm));
0135     m_vy.push_back(Acts::clampValue<float>(particle.fourPosition().y() /
0136                                            Acts::UnitConstants::mm));
0137     m_vz.push_back(Acts::clampValue<float>(particle.fourPosition().z() /
0138                                            Acts::UnitConstants::mm));
0139     m_vt.push_back(Acts::clampValue<float>(particle.fourPosition().w() /
0140                                            Acts::UnitConstants::mm));
0141
0142     // particle constants
0143     if (!std::isfinite(particle.mass()) || !std::isfinite(particle.charge())) {
0144       ACTS_WARNING("Particle mass or charge is not finite, can't write it");
0145     }
0146
0147     m_m.push_back(
0148         Acts::clampValue<float>(particle.mass() / Acts::UnitConstants::GeV));
0149     m_q.push_back(
0150         Acts::clampValue<float>(particle.charge() / Acts::UnitConstants::e));
0151     // decoded barcode components
0152     m_vertexPrimary.push_back(particle.particleId().vertexPrimary());
0153     m_vertexSecondary.push_back(particle.particleId().vertexSecondary());
0154     m_particle.push_back(particle.particleId().particle());
0155     m_generation.push_back(particle.particleId().generation());
0156     m_subParticle.push_back(particle.particleId().subParticle());
0157
0158     m_eLoss.push_back(Acts::clampValue<float>(particle.energyLoss() /
0159                                               Acts::UnitConstants::GeV));
0160     m_pathInX0.push_back(
0161         Acts::clampValue<float>(particle.pathInX0() / Acts::UnitConstants::mm));
0162     m_pathInL0.push_back(
0163         Acts::clampValue<float>(particle.pathInL0() / Acts::UnitConstants::mm));
0164     m_numberOfHits.push_back(particle.numberOfHits());
0165     m_outcome.push_back(static_cast<std::uint32_t>(particle.outcome()));
0166
0167     // momentum
0168     const auto p = particle.absoluteMomentum() / Acts::UnitConstants::GeV;
0169     m_p.push_back(Acts::clampValue<float>(p));
0170     m_px.push_back(Acts::clampValue<float>(p * particle.direction().x()));
0171     m_py.push_back(Acts::clampValue<float>(p * particle.direction().y()));
0172     m_pz.push_back(Acts::clampValue<float>(p * particle.direction().z()));
0173     // derived kinematic quantities
0174     m_eta.push_back(Acts::clampValue<float>(
0175         Acts::VectorHelpers::eta(particle.direction())));
0176     m_pt.push_back(Acts::clampValue<float>(
0177         p * Acts::VectorHelpers::perp(particle.direction())));
0178     m_phi.push_back(Acts::clampValue<float>(
0179         Acts::VectorHelpers::phi(particle.direction())));
0180     m_theta.push_back(Acts::clampValue<float>(
0181         Acts::VectorHelpers::theta(particle.direction())));
0182     m_qop.push_back(Acts::clampValue<float>(
0183         particle.qOverP() * Acts::UnitConstants::GeV / Acts::UnitConstants::e));
0184
0185     if (!m_cfg.writeHelixParameters) {
0186       // done with this particle
0187       continue;
0188     }
0189
0190     // Perigee surface at configured reference point
0191     auto pSurface =
0192         Acts::Surface::makeShared<Acts::PerigeeSurface>(m_cfg.referencePoint);
0193
0194     // Start from truth curvilinear parameters (direction, q/p)
0195     const Acts::Vector3 startDir = particle.direction();  // unit vector
0196     const auto qOverP = particle.qOverP();                // ACTS units
0197     // The perigee surface is not aligned, so either context works
0198     auto intersection =
0199         pSurface
0200             ->intersect(ctx.recoGeoContext, particle.position(), startDir,
0201                         Acts::BoundaryTolerance::Infinite())
0202             .closest();
0203
0204     // Neutral particles have no helix -> linearly extrapolate to perigee
0205     if (particle.charge() == 0) {
0206       ACTS_WARNING(
0207           "Particle has zero charge, linearly extrapolating to perigee");
0208       // Initialize the truth particle info
0209       auto perigeeD0 = NaNfloat;
0210       auto perigeeZ0 = NaNfloat;
0211
0212       const auto position = intersection.position();
0213
0214       // get the truth perigee parameter
0215       auto lpResult =
0216           pSurface->globalToLocal(ctx.recoGeoContext, position, startDir);
0217       if (lpResult.ok()) {
0218         perigeeD0 = lpResult.value()[Acts::BoundIndices::eBoundLoc0];
0219         perigeeZ0 = lpResult.value()[Acts::BoundIndices::eBoundLoc1];
0220       } else {
0221         ACTS_ERROR("Global to local transformation did not succeed.");
0222       }
0223       // truth parameters at perigee are the same as at production vertex
0224       m_perigeePhi.push_back(Acts::clampValue<float>(particle.phi()));
0225       m_perigeeTheta.push_back(Acts::clampValue<float>(particle.theta()));
0226       m_perigeeQop.push_back(Acts::clampValue<float>(
0227           qOverP * Acts::UnitConstants::GeV / Acts::UnitConstants::e));
0228       m_perigeeP.push_back(Acts::clampValue<float>(particle.absoluteMomentum() /
0229                                                    Acts::UnitConstants::GeV));
0230       m_perigeePx.push_back(Acts::clampValue<float>(m_p.back() * startDir.x()));
0231       m_perigeePy.push_back(Acts::clampValue<float>(m_p.back() * startDir.y()));
0232       m_perigeePz.push_back(Acts::clampValue<float>(m_p.back() * startDir.z()));
0233       m_perigeeEta.push_back(Acts::clampValue<float>(
0234           Acts::VectorHelpers::eta(particle.direction())));
0235       m_perigeePt.push_back(Acts::clampValue<float>(
0236           m_p.back() * Acts::VectorHelpers::perp(particle.direction())));
0237
0238       // Push the extrapolated parameters
0239       m_perigeeD0.push_back(
0240           Acts::clampValue<float>(perigeeD0 / Acts::UnitConstants::mm));
0241       m_perigeeZ0.push_back(
0242           Acts::clampValue<float>(perigeeZ0 / Acts::UnitConstants::mm));
0243       continue;
0244     }
0245
0246     // Charged particles: propagate helix to perigee
0247     // Build a propagator and propagate the *truth parameters* to the
0248     // perigee Stepper + propagator
0249     using Stepper = Acts::SympyStepper;
0250     Stepper stepper(m_cfg.bField);
0251     using PropagatorT = Acts::Propagator<Stepper>;
0252     auto propagator = std::make_shared<PropagatorT>(stepper);
0253
0254     Acts::BoundTrackParameters startParams =
0255         Acts::BoundTrackParameters::createCurvilinear(
0256             particle.fourPosition(), startDir, qOverP, std::nullopt,
0257             Acts::ParticleHypothesis::pion());
0258
0259     // Propagation options (need event contexts)
0260     using PropOptions = PropagatorT::Options<>;
0261     PropOptions pOptions(ctx.recoGeoContext, ctx.magFieldContext);
0262
0263     // Choose propagation direction based on the closest intersection
0264     pOptions.direction =
0265         Acts::Direction::fromScalarZeroAsPositive(intersection.pathLength());
0266
0267     // Do the propagation to the perigee surface
0268     auto propRes = propagator->propagate(startParams, *pSurface, pOptions);
0269     if (!propRes.ok() || !propRes->endParameters.has_value()) {
0270       ACTS_ERROR("Propagation to perigee surface failed.");
0271       m_perigeePhi.push_back(NaNfloat);
0272       m_perigeeTheta.push_back(NaNfloat);
0273       m_perigeeQop.push_back(NaNfloat);
0274       m_perigeeD0.push_back(NaNfloat);
0275       m_perigeeZ0.push_back(NaNfloat);
0276       m_perigeeP.push_back(NaNfloat);
0277       m_perigeePx.push_back(NaNfloat);
0278       m_perigeePy.push_back(NaNfloat);
0279       m_perigeePz.push_back(NaNfloat);
0280       m_perigeeEta.push_back(NaNfloat);
0281       m_perigeePt.push_back(NaNfloat);
0282       continue;
0283     }
0284     const Acts::BoundTrackParameters& atPerigee = *propRes->endParameters;
0285
0286     // By construction, atPerigee is *bound on the perigee surface*.
0287     // Its parameter vector is [loc0, loc1, phi, theta, q/p, t]
0288     const auto& perigee_pars = atPerigee.parameters();
0289
0290     const auto perigeeD0 = perigee_pars[Acts::BoundIndices::eBoundLoc0];
0291     const auto perigeeZ0 = perigee_pars[Acts::BoundIndices::eBoundLoc1];
0292
0293     // truth phi;theta;q/p *at the perigee*,
0294     const auto perigeePhi = perigee_pars[Acts::BoundIndices::eBoundPhi];
0295     const auto perigeeTheta = perigee_pars[Acts::BoundIndices::eBoundTheta];
0296     const auto perigeeQop = perigee_pars[Acts::BoundIndices::eBoundQOverP];
0297
0298     m_perigeePhi.push_back(Acts::clampValue<float>(perigeePhi));
0299     m_perigeeTheta.push_back(Acts::clampValue<float>(perigeeTheta));
0300     m_perigeeQop.push_back(Acts::clampValue<float>(
0301         perigeeQop * Acts::UnitConstants::GeV / Acts::UnitConstants::e));
0302     // update p, px, py, pz, eta, pt
0303     const auto perigeeP =
0304         atPerigee.absoluteMomentum() / Acts::UnitConstants::GeV;
0305     m_perigeeP.push_back(Acts::clampValue<float>(perigeeP));
0306     const auto dir = atPerigee.direction();
0307     m_perigeePx.push_back(Acts::clampValue<float>(perigeeP * dir.x()));
0308     m_perigeePy.push_back(Acts::clampValue<float>(perigeeP * dir.y()));
0309     m_perigeePz.push_back(Acts::clampValue<float>(perigeeP * dir.z()));
0310     m_perigeeEta.push_back(Acts::clampValue<float>(
0311         Acts::VectorHelpers::eta(atPerigee.direction())));
0312     m_perigeePt.push_back(Acts::clampValue<float>(
0313         perigeeP * Acts::VectorHelpers::perp(atPerigee.direction())));
0314
0315     // Push the perigee parameters
0316     m_perigeeD0.push_back(
0317         Acts::clampValue<float>(perigeeD0 / Acts::UnitConstants::mm));
0318     m_perigeeZ0.push_back(
0319         Acts::clampValue<float>(perigeeZ0 / Acts::UnitConstants::mm));
0320   }
0321
0322   m_outputTree->Fill();
0323
0324   m_particleHash.clear();
0325   m_particleType.clear();
0326   m_process.clear();
0327   m_vx.clear();
0328   m_vy.clear();
0329   m_vz.clear();
0330   m_vt.clear();
0331   m_p.clear();
0332   m_px.clear();
0333   m_py.clear();
0334   m_pz.clear();
0335   m_m.clear();
0336   m_q.clear();
0337   m_eta.clear();
0338   m_phi.clear();
0339   m_pt.clear();
0340   m_theta.clear();
0341   m_qop.clear();
0342   m_vertexPrimary.clear();
0343   m_vertexSecondary.clear();
0344   m_particle.clear();
0345   m_generation.clear();
0346   m_subParticle.clear();
0347   m_eLoss.clear();
0348   m_numberOfHits.clear();
0349   m_outcome.clear();
0350   m_pathInX0.clear();
0351   m_pathInL0.clear();
0352
0353   if (m_cfg.writeHelixParameters) {
0354     m_perigeeD0.clear();
0355     m_perigeeZ0.clear();
0356     m_perigeePhi.clear();
0357     m_perigeeTheta.clear();
0358     m_perigeeQop.clear();
0359     m_perigeeP.clear();
0360     m_perigeePx.clear();
0361     m_perigeePy.clear();
0362     m_perigeePz.clear();
0363     m_perigeeEta.clear();
0364     m_perigeePt.clear();
0365   }
0366
0367   m_origParticleIdx.clear();
0368   m_hfOrigin.clear();
0369
0370   return ProcessCode::SUCCESS;
0371 }
0372
0373 }  // namespace ActsExamples