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   /**
0028    * @brief Produce a list of primary vertex candidates
0029    *
0030    * @param rcvtx  - input collection of all vertex candidates
0031    * @return edm4eic::VertexCollection
0032    */
0033   void PrimaryVertices::process(const PrimaryVertices::Input& input,
0034                 const PrimaryVertices::Output& output) const {
0035     const auto [rcvtx] = input;
0036     auto [out_primary_vertices] = output;
0037
0038     // this multimap will store intermediate results
0039     // so that we can sort them before filling output
0040     // collection
0041     std::multimap<int, edm4eic::Vertex, std::greater<int>> primaryVertexMap;
0042
0043     // our output collection of primary vertex
0044     // ordered by N_trk = associatedParticle array size
0045     out_primary_vertices->setSubsetCollection();
0046
0047     trace("We have {} candidate vertices", rcvtx->size());
0048
0049     for (const auto& vtx: *rcvtx) {
0050       const auto &v = vtx.getPosition();
0051
0052       // some basic vertex selection
0053       if (sqrt( v.x*v.x + v.y*v.y ) / edm4eic::unit::mm > m_cfg.maxVr ||
0054           fabs( v.z ) / edm4eic::unit::mm > m_cfg.maxVz )
0055         continue;
0056
0057       if (vtx.getChi2() > m_cfg.maxChi2) continue;
0058
0059       int N_trk = vtx.getAssociatedParticles().size();
0060       trace( "\t N_trk = {}", N_trk );
0061       primaryVertexMap.insert({N_trk, vtx});
0062     } // vertex loop
0063
0064     bool first = true;
0065     // map defined with std::greater<> will be iterated in descending order by the key
0066     for (auto [N_trk, vertex] : primaryVertexMap) {
0067       // Do not save primary candidates that
0068       // are not within range
0069       if ( N_trk > m_cfg.maxNtrk
0070         || N_trk < m_cfg.minNtrk ){
0071         continue;
0072       }
0073
0074       // For logging and development
0075       // report the highest N_trk candidate chosen
0076       if (first) {
0077         trace("Max N_trk Candidate:");
0078         trace("\t N_trk = {}", N_trk);
0079         trace("\t Primary vertex has xyz=( {}, {}, {} )", vertex.getPosition().x, vertex.getPosition().y, vertex.getPosition().z);
0080         first = false;
0081       }
0082       out_primary_vertices->push_back(vertex);
0083     } // loop on primaryVertexMap
0084   }
0085 } // namespace eicrecon