src/components/ImagingPixelDataSorter.cpp

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2022 Chao Peng, Chao, Sylvester Joosten
0003
0004 /*
0005  *  A hits-level data sorter to prepare dataset for machine learning
0006  *
0007  *  Author: Chao Peng (ANL), 05/04/2021
0008  */
0009 #include <algorithm>
0010 #include <bitset>
0011 #include <unordered_map>
0012
0013 #include "Gaudi/Property.h"
0014 #include "GaudiAlg/GaudiAlgorithm.h"
0015 #include "GaudiAlg/GaudiTool.h"
0016 #include "GaudiAlg/Transformer.h"
0017 #include "GaudiKernel/PhysicalConstants.h"
0018 #include "GaudiKernel/RndmGenerators.h"
0019 #include "GaudiKernel/ToolHandle.h"
0020
0021 #include "DDRec/CellIDPositionConverter.h"
0022 #include "DDRec/Surface.h"
0023 #include "DDRec/SurfaceManager.h"
0024
0025 #include <k4FWCore/DataHandle.h>
0026 #include <k4Interface/IGeoSvc.h>
0027
0028 // Event Model related classes
0029 #include <edm4hep/utils/vector_utils.h>
0030 #include "edm4eic/CalorimeterHit.h"
0031 #include "edm4eic/CalorimeterHitCollection.h"
0032
0033 using namespace Gaudi::Units;
0034
0035 namespace Jug::Reco {
0036
0037   /** Hits sorter for ML algorithm input.
0038    *
0039    * A hits-level data sorter to prepare dataset for machine learning
0040    * It sorts the hits by layer and energy with defined sizes (max number of layers and max number of hits per layer)
0041    * Hits are sorted by energy in a descending order.
0042    * Out-of-range hits will be discarded and empty slots will be padded with zeros
0043    *
0044    * \ingroup reco
0045    */
0046   class ImagingPixelDataSorter : public GaudiAlgorithm {
0047   private:
0048     Gaudi::Property<int>                        m_nLayers{this, "numberOfLayers", 9};
0049     Gaudi::Property<int>                        m_nHits{this, "numberOfHits", 50};
0050     DataHandle<edm4eic::CalorimeterHitCollection>   m_inputHitCollection{"inputHitCollection",
0051                                                                      Gaudi::DataHandle::Reader, this};
0052     DataHandle<edm4eic::CalorimeterHitCollection>   m_outputHitCollection{"outputHitCollection",
0053                                                                       Gaudi::DataHandle::Writer, this};
0054
0055   public:
0056     ImagingPixelDataSorter(const std::string& name, ISvcLocator* svcLoc)
0057       : GaudiAlgorithm(name, svcLoc)
0058     {
0059       declareProperty("inputHitCollection", m_inputHitCollection, "");
0060       declareProperty("outputHitCollection", m_outputHitCollection, "");
0061     }
0062
0063     StatusCode initialize() override
0064     {
0065       if (GaudiAlgorithm::initialize().isFailure()) {
0066         return StatusCode::FAILURE;
0067       }
0068
0069       return StatusCode::SUCCESS;
0070     }
0071
0072     StatusCode execute() override
0073     {
0074       // input collections
0075       const auto& hits = *m_inputHitCollection.get();
0076       // Create output collections
0077       auto& mhits = *m_outputHitCollection.createAndPut();
0078
0079       // group the hits by layer
0080       std::vector<std::vector<edm4eic::CalorimeterHit>> layer_hits;
0081       layer_hits.resize(m_nLayers);
0082       for (const auto& h : hits) {
0083         auto k = h.getLayer();
0084         if ((int)k < m_nLayers) {
0085           layer_hits[k].push_back(h);
0086         }
0087       }
0088
0089       // sort by energy
0090       for (auto &layer : layer_hits) {
0091         std::sort(layer.begin(), layer.end(),
0092           [] (const edm4eic::CalorimeterHit &h1, const edm4eic::CalorimeterHit &h2) {
0093             return h1.getEnergy() > h2.getEnergy();
0094           });
0095       }
0096
0097       // fill-in the output
0098       for (size_t k = 0; k < layer_hits.size(); ++k) {
0099         auto &layer = layer_hits[k];
0100         for (size_t i = 0; i < (size_t) m_nHits; ++i) {
0101           // pad zeros if no hits
0102           if (i >= layer.size()) {
0103             auto h = mhits.create();
0104             h.setLayer((int)k);
0105             h.setEnergy(0.);
0106           } else {
0107             auto h = layer[i].clone();
0108             mhits.push_back(h);
0109           }
0110         }
0111       }
0112
0113       return StatusCode::SUCCESS;
0114     }
0115
0116   }; // class ImagingPixelDataSorter
0117
0118   // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
0119   DECLARE_COMPONENT(ImagingPixelDataSorter)
0120
0121 } // namespace Jug::Reco