Back to home page

EIC code displayed by LXR

 
 

    


Warning, file /juggler/JugReco/src/components/ImagingPixelDataCombiner.cpp was not indexed or was modified since last indexation (in which case cross-reference links may be missing, inaccurate or erroneous).

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2022 Chao Peng, Sylvester Joosten, Whitney Armstrong
0003 
0004 /*
0005  *  A hits-level data combiner to combine two datasets into one for machine learning
0006  *
0007  *  Author: Chao Peng (ANL), 05/04/2021
0008  */
0009 #include <algorithm>
0010 #include <bitset>
0011 #include <fmt/format.h>
0012 #include <unordered_map>
0013 
0014 #include "Gaudi/Property.h"
0015 #include "Gaudi/Algorithm.h"
0016 #include "GaudiKernel/PhysicalConstants.h"
0017 #include "GaudiKernel/RndmGenerators.h"
0018 #include "GaudiKernel/ToolHandle.h"
0019 
0020 #include "DDRec/CellIDPositionConverter.h"
0021 #include "DDRec/Surface.h"
0022 #include "DDRec/SurfaceManager.h"
0023 
0024 #include <k4FWCore/DataHandle.h>
0025 
0026 // Event Model related classes
0027 #include "edm4eic/CalorimeterHitCollection.h"
0028 #include "edm4hep/utils/vector_utils.h"
0029 
0030 using namespace Gaudi::Units;
0031 
0032 namespace Jug::Reco {
0033 
0034 /** Hits combiner for ML algorithm input.
0035  *
0036  * A hits-level data combiner to combine two datasets into one for machine learning
0037  * It accepts inputs from data sorter that hits are sorted by layers
0038  * Two different datasets will be combined together following specified rules in handling the layers
0039  * Supported rules: concatenate, interlayer
0040  *
0041  * \ingroup reco
0042  */
0043 class ImagingPixelDataCombiner : public Gaudi::Algorithm {
0044 private:
0045   Gaudi::Property<int> m_layerIncrement{this, "layerIncrement", 0};
0046   Gaudi::Property<std::string> m_rule{this, "rule", "concatenate"};
0047   mutable DataHandle<edm4eic::CalorimeterHitCollection> m_inputHits1{"inputHits1", Gaudi::DataHandle::Reader, this};
0048   mutable DataHandle<edm4eic::CalorimeterHitCollection> m_inputHits2{"inputHits2", Gaudi::DataHandle::Reader, this};
0049   mutable DataHandle<edm4eic::CalorimeterHitCollection> m_outputHits{"outputHits", Gaudi::DataHandle::Writer, this};
0050   std::vector<std::string> supported_rules{"concatenate", "interlayer"};
0051 
0052 public:
0053   ImagingPixelDataCombiner(const std::string& name, ISvcLocator* svcLoc)
0054       : Gaudi::Algorithm(name, svcLoc) {
0055     declareProperty("inputHits1", m_inputHits1, "");
0056     declareProperty("inputHits2", m_inputHits2, "");
0057     declareProperty("outputHits", m_outputHits, "");
0058   }
0059 
0060   StatusCode initialize() override {
0061     if (Gaudi::Algorithm::initialize().isFailure()) {
0062       return StatusCode::FAILURE;
0063     }
0064 
0065     if (std::find(supported_rules.begin(), supported_rules.end(), m_rule.value()) == supported_rules.end()) {
0066       error() << fmt::format("unsupported rule: {}, please choose one from [{}]", m_rule.value(),
0067                              fmt::join(supported_rules, ", "))
0068               << endmsg;
0069       return StatusCode::FAILURE;
0070     }
0071 
0072     return StatusCode::SUCCESS;
0073   }
0074 
0075   StatusCode execute(const EventContext&) const override {
0076     // input collections
0077     const auto* const hits1 = m_inputHits1.get();
0078     const auto* const hits2 = m_inputHits2.get();
0079     std::vector<const edm4eic::CalorimeterHitCollection*> inputs{hits1, hits2};
0080     // Create output collections
0081     auto* mhits = m_outputHits.createAndPut();
0082 
0083     // concatenate
0084     if (m_rule.value() == supported_rules[0]) {
0085       for (int i = 0; i < (int)inputs.size(); ++i) {
0086         const auto* const coll = inputs[i];
0087         for (auto hit : *coll) {
0088           edm4eic::CalorimeterHit h2{
0089               hit.getCellID(),    hit.getEnergy(),   hit.getEnergyError(), hit.getTime(),
0090               hit.getTimeError(), hit.getPosition(), hit.getDimension(),   hit.getLayer() + m_layerIncrement * i,
0091               hit.getSector(),    hit.getLocal(),
0092           };
0093           mhits->push_back(h2);
0094         }
0095       }
0096       // interlayer
0097       // @NOTE: it assumes the input hits are sorted by layers
0098     } else if (m_rule.value() == supported_rules[1]) {
0099       std::vector<int> indices{0, 0};
0100       int curr_coll   = 0;
0101       bool init_layer = false;
0102       int curr_layer  = 0;
0103       // int curr_ihit = 0;
0104       while (indices[0] < (int)hits1->size() || indices[1] < (int)hits2->size()) {
0105         // cyclic index
0106         if (curr_coll >= (int)inputs.size()) {
0107           curr_coll -= (int)inputs.size();
0108         }
0109 
0110         // merge hits
0111         int& i                 = indices[curr_coll];
0112         const auto* const coll = inputs[curr_coll];
0113 
0114         // reach this collection's end
0115         if (i >= (int)coll->size()) {
0116           curr_coll++;
0117           init_layer = false;
0118           // curr_ihit = 0;
0119           // info() << "collection end" << endmsg;
0120           continue;
0121         }
0122 
0123         auto hit = (*coll)[i];
0124         if (!init_layer) {
0125           curr_layer = hit.getLayer();
0126           init_layer = true;
0127         }
0128 
0129         // reach this layer's end
0130         if (curr_layer != hit.getLayer()) {
0131           curr_coll++;
0132           init_layer = false;
0133           // curr_ihit = 0;
0134           // info() << "layer end : " << curr_layer << " != " << hit.getLayer() << endmsg;
0135           continue;
0136         }
0137 
0138         // push hit, increment of index
0139         edm4eic::CalorimeterHit h2{
0140             hit.getCellID(),    hit.getEnergy(),   hit.getEnergyError(), hit.getTime(),
0141             hit.getTimeError(), hit.getPosition(), hit.getDimension(),   hit.getLayer() + m_layerIncrement * curr_coll,
0142             hit.getSector(),    hit.getLocal()};
0143         mhits->push_back(h2);
0144         i++;
0145         // info() << curr_coll << ": " << curr_ihit ++ << endmsg;
0146       }
0147       // info() << hits1->size() << ", " << hits2->size() << endmsg;
0148     }
0149 
0150     return StatusCode::SUCCESS;
0151   }
0152 
0153 }; // class ImagingPixelDataCombiner
0154 
0155 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
0156 DECLARE_COMPONENT(ImagingPixelDataCombiner)
0157 
0158 } // namespace Jug::Reco