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