src/components/TruthClustering.cpp

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2022 Sylvester Joosten, Whitney Armstrong, Wouter Deconinck
0003
0004 #include <algorithm>
0005
0006 #include "Gaudi/Property.h"
0007 #include "GaudiAlg/GaudiAlgorithm.h"
0008 #include "GaudiAlg/GaudiTool.h"
0009 #include "GaudiAlg/Transformer.h"
0010 #include "GaudiKernel/PhysicalConstants.h"
0011 #include "GaudiKernel/RndmGenerators.h"
0012 #include "GaudiKernel/ToolHandle.h"
0013
0014 #include "DDRec/CellIDPositionConverter.h"
0015 #include "DDRec/Surface.h"
0016 #include "DDRec/SurfaceManager.h"
0017
0018 #include <k4FWCore/DataHandle.h>
0019 #include <k4Interface/IGeoSvc.h>
0020
0021 // Event Model related classes
0022 #include "edm4hep/MCParticle.h"
0023 #include "edm4hep/SimCalorimeterHitCollection.h"
0024 #include "edm4eic/CalorimeterHitCollection.h"
0025 #include "edm4eic/ClusterCollection.h"
0026 #include "edm4eic/ProtoClusterCollection.h"
0027
0028 using namespace Gaudi::Units;
0029
0030 namespace Jug::Fast {
0031
0032 /** Truth clustering algorithm.
0033  *
0034  * \ingroup reco
0035  */
0036 class TruthClustering : public GaudiAlgorithm {
0037 private:
0038   DataHandle<edm4eic::CalorimeterHitCollection> m_inputHits{"inputHits", Gaudi::DataHandle::Reader, this};
0039   DataHandle<edm4hep::SimCalorimeterHitCollection> m_mcHits{"mcHits", Gaudi::DataHandle::Reader, this};
0040   DataHandle<edm4eic::ProtoClusterCollection> m_outputProtoClusters{"outputProtoClusters", Gaudi::DataHandle::Writer, this};
0041
0042 public:
0043   TruthClustering(const std::string& name, ISvcLocator* svcLoc)
0044       : GaudiAlgorithm(name, svcLoc) {
0045     declareProperty("inputHits", m_inputHits, "Input calorimeter reco hits");
0046     declareProperty("mcHits", m_mcHits, "Input truth hits");
0047     declareProperty("outputProtoClusters", m_outputProtoClusters, "Output proto clusters");
0048   }
0049
0050   StatusCode initialize() override {
0051     if (GaudiAlgorithm::initialize().isFailure()) {
0052       return StatusCode::FAILURE;
0053     }
0054     return StatusCode::SUCCESS;
0055   }
0056
0057   StatusCode execute() override {
0058     // input collections
0059     const auto& hits = *m_inputHits.get();
0060     const auto& mc   = *m_mcHits.get();
0061     // Create output collections
0062     auto& proto = *m_outputProtoClusters.createAndPut();
0063
0064     // Map mc track ID to protoCluster index
0065     std::map<int32_t, int32_t> protoIndex;
0066
0067     // Loop over al calorimeter hits and sort per mcparticle
0068     for (const auto& hit : hits) {
0069       const auto& mcHit     = mc[hit.getObjectID().index];
0070       const auto& trackID   = mcHit.getContributions(0).getParticle().getObjectID().index;
0071       // Create a new protocluster if we don't have one for this trackID
0072       if (protoIndex.count(trackID) == 0) {
0073         auto pcl = proto.create();
0074         protoIndex[trackID] = proto.size() - 1;
0075       }
0076       // Add hit to the appropriate protocluster
0077       proto[protoIndex[trackID]].addToHits(hit);
0078       proto[protoIndex[trackID]].addToWeights(1);
0079     }
0080     return StatusCode::SUCCESS;
0081   }
0082 };
0083
0084 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
0085 DECLARE_COMPONENT(TruthClustering)
0086
0087 } // namespace Jug::Fast