truth/todo/TruthClustering.h

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 "JugBase/DataHandle.h"
0019 #include "JugBase/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 #include "edm4eic/RawCalorimeterHitCollection.h"
0028
0029 using namespace Gaudi::Units;
0030
0031 namespace Jug::Fast {
0032
0033 /** Truth clustering algorithm.
0034  *
0035  * \ingroup reco
0036  */
0037 class TruthClustering : public GaudiAlgorithm {
0038 private:
0039   DataHandle<edm4eic::CalorimeterHitCollection> m_inputHits{"inputHits", Gaudi::DataHandle::Reader, this};
0040   DataHandle<edm4hep::SimCalorimeterHitCollection> m_mcHits{"mcHits", Gaudi::DataHandle::Reader, this};
0041   DataHandle<edm4eic::ProtoClusterCollection> m_outputProtoClusters{"outputProtoClusters", Gaudi::DataHandle::Writer, this};
0042
0043 public:
0044   TruthClustering(const std::string& name, ISvcLocator* svcLoc)
0045       : GaudiAlgorithm(name, svcLoc) {
0046     declareProperty("inputHits", m_inputHits, "Input calorimeter reco hits");
0047     declareProperty("mcHits", m_mcHits, "Input truth hits");
0048     declareProperty("outputProtoClusters", m_outputProtoClusters, "Output proto clusters");
0049   }
0050
0051   StatusCode initialize() override {
0052     if (GaudiAlgorithm::initialize().isFailure()) {
0053       return StatusCode::FAILURE;
0054     }
0055     return StatusCode::SUCCESS;
0056   }
0057
0058   StatusCode execute() override {
0059     // input collections
0060     const auto& hits = *m_inputHits.get();
0061     const auto& mc   = *m_mcHits.get();
0062     // Create output collections
0063     auto& proto = *m_outputProtoClusters.createAndPut();
0064
0065     // Map mc track ID to protoCluster index
0066     std::map<int32_t, int32_t> protoIndex;
0067
0068     // Loop over al calorimeter hits and sort per mcparticle
0069     for (const auto& hit : hits) {
0070       const auto& mcHit     = mc[hit.getObjectID().index];
0071       const auto& trackID   = mcHit.getContributions(0).getParticle().id();
0072       // Create a new protocluster if we don't have one for this trackID
0073       if (protoIndex.count(trackID) == 0) {
0074         auto pcl = proto.create();
0075         protoIndex[trackID] = proto.size() - 1;
0076       }
0077       // Add hit to the appropriate protocluster
0078       proto[protoIndex[trackID]].addToHits(hit);
0079       proto[protoIndex[trackID]].addToWeights(1);
0080     }
0081     return StatusCode::SUCCESS;
0082   }
0083 };
0084
0085 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
0086 DECLARE_COMPONENT(TruthClustering)
0087
0088 } // namespace Jug::Fast