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