algorithms/reco/PrimaryVertices.cc

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2024 Daniel Brandenburg, Xin Dong
0003
0004 #include <edm4eic/ReconstructedParticle.h>
0005 #include <edm4eic/VertexCollection.h>
0006 #include <edm4eic/unit_system.h>
0007 #include <edm4hep/Vector4f.h>
0008 #include <fmt/core.h>
0009 #include <podio/RelationRange.h>
0010 #include <cmath>
0011 #include <functional>
0012 #include <gsl/pointers>
0013 #include <map>
0014
0015 #include "algorithms/reco/PrimaryVertices.h"
0016 #include "algorithms/reco/PrimaryVerticesConfig.h"
0017
0018 namespace eicrecon {
0019
0020 /**
0021    * @brief Initialize the PrimaryVertices Algorithm
0022    *
0023    */
0024 void PrimaryVertices::init() {}
0025
0026 /**
0027    * @brief Produce a list of primary vertex candidates
0028    *
0029    * @param rcvtx  - input collection of all vertex candidates
0030    * @return edm4eic::VertexCollection
0031    */
0032 void PrimaryVertices::process(const PrimaryVertices::Input& input,
0033                               const PrimaryVertices::Output& output) const {
0034   const auto [rcvtx]          = input;
0035   auto [out_primary_vertices] = output;
0036
0037   // this multimap will store intermediate results
0038   // so that we can sort them before filling output
0039   // collection
0040   std::multimap<int, edm4eic::Vertex, std::greater<>> primaryVertexMap;
0041
0042   // our output collection of primary vertex
0043   // ordered by N_trk = associatedParticle array size
0044   out_primary_vertices->setSubsetCollection();
0045
0046   trace("We have {} candidate vertices", rcvtx->size());
0047
0048   for (const auto& vtx : *rcvtx) {
0049     const auto& v = vtx.getPosition();
0050
0051     // some basic vertex selection
0052     if (sqrt(v.x * v.x + v.y * v.y) / edm4eic::unit::mm > m_cfg.maxVr ||
0053         std::abs(v.z) / edm4eic::unit::mm > m_cfg.maxVz) {
0054       continue;
0055     }
0056
0057     if (vtx.getChi2() > m_cfg.maxChi2) {
0058       continue;
0059     }
0060
0061     int N_trk = vtx.getAssociatedParticles().size();
0062     trace("\t N_trk = {}", N_trk);
0063     primaryVertexMap.insert({N_trk, vtx});
0064   } // vertex loop
0065
0066   bool first = true;
0067   // map defined with std::greater<> will be iterated in descending order by the key
0068   for (auto [N_trk, vertex] : primaryVertexMap) {
0069     // Do not save primary candidates that
0070     // are not within range
0071     if (N_trk > m_cfg.maxNtrk || N_trk < m_cfg.minNtrk) {
0072       continue;
0073     }
0074
0075     // For logging and development
0076     // report the highest N_trk candidate chosen
0077     if (first) {
0078       trace("Max N_trk Candidate:");
0079       trace("\t N_trk = {}", N_trk);
0080       trace("\t Primary vertex has xyz=( {}, {}, {} )", vertex.getPosition().x,
0081             vertex.getPosition().y, vertex.getPosition().z);
0082       first = false;
0083     }
0084     out_primary_vertices->push_back(vertex);
0085   } // loop on primaryVertexMap
0086 }
0087 } // namespace eicrecon