calorimetry/todo/ImagingPixelDataSorter.h

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