Back to home page

EIC code displayed by LXR

 
 

    


Warning, file /juggler/JugReco/src/components/ImagingPixelMerger.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, Wouter Deconinck, Whitney Armstrong
0003 
0004 /*
0005  *  A hits merger for ecal barrel to prepare dataset for machine learning
0006  *
0007  *  Author: Chao Peng (ANL), 05/04/2021
0008  *
0009  *  SJJ: @TODO: this should really be a clustering algorithm, as it doesn't give us
0010  *              fully consistent hits anymore (e.g. cellID, local position) (meaning it's
0011  *              not generically useful as a reconstruction algorithm.
0012  */
0013 #include <algorithm>
0014 #include <bitset>
0015 #include <unordered_map>
0016 
0017 #include "Gaudi/Property.h"
0018 #include "Gaudi/Algorithm.h"
0019 #include "GaudiKernel/PhysicalConstants.h"
0020 #include "GaudiKernel/RndmGenerators.h"
0021 #include "GaudiKernel/ToolHandle.h"
0022 
0023 #include "DDRec/CellIDPositionConverter.h"
0024 #include "DDRec/Surface.h"
0025 #include "DDRec/SurfaceManager.h"
0026 
0027 #include <k4FWCore/DataHandle.h>
0028 #include <k4Interface/IGeoSvc.h>
0029 #include "JugBase/Utilities/Utils.hpp"
0030 
0031 // Event Model related classes
0032 #include "edm4eic/CalorimeterHitCollection.h"
0033 #include <edm4hep/utils/vector_utils.h>
0034 
0035 using namespace Gaudi::Units;
0036 
0037 struct PairHashFunction {
0038   template <class T1, class T2> std::size_t operator()(const std::pair<T1, T2>& pair) const {
0039     return std::hash<T1>()(pair.first) ^ std::hash<T2>()(pair.second);
0040   }
0041 };
0042 
0043 namespace Jug::Reco {
0044 
0045 /** Hits merger for ML algorithm input.
0046  *
0047  * A hits merger to prepare dataset for machine learning
0048  * It merges hits with a defined grid size.
0049  * Merged hits will be relocated to the grid center and the energies will be summed.
0050  *
0051  * \ingroup reco
0052  */
0053 class ImagingPixelMerger : public Gaudi::Algorithm {
0054 private:
0055   Gaudi::Property<float> m_etaSize{this, "etaSize", 0.001};
0056   Gaudi::Property<float> m_phiSize{this, "phiSize", 0.001};
0057   mutable DataHandle<edm4eic::CalorimeterHitCollection> m_inputHits{"inputHits", Gaudi::DataHandle::Reader, this};
0058   mutable DataHandle<edm4eic::CalorimeterHitCollection> m_outputHits{"outputHits", Gaudi::DataHandle::Writer, this};
0059 
0060 public:
0061   ImagingPixelMerger(const std::string& name, ISvcLocator* svcLoc) : Gaudi::Algorithm(name, svcLoc) {
0062     declareProperty("inputHits", m_inputHits, "");
0063     declareProperty("outputHits", m_outputHits, "");
0064   }
0065 
0066   StatusCode initialize() override {
0067     if (Gaudi::Algorithm::initialize().isFailure()) {
0068       return StatusCode::FAILURE;
0069     }
0070 
0071     return StatusCode::SUCCESS;
0072   }
0073 
0074   StatusCode execute(const EventContext&) const override {
0075     // input collections
0076     const auto& hits = *m_inputHits.get();
0077     // Create output collections
0078     auto& ohits = *m_outputHits.createAndPut();
0079 
0080     // @TODO: add timing information
0081     // group the hits by grid per layer
0082     struct GridData {
0083       unsigned nHits;
0084       float rc;
0085       float energy;
0086       float energyError;
0087       float time;
0088       float timeError;
0089       int sector; // sector associated with one of the merged hits
0090     };
0091     // @TODO: remove this hard-coded value
0092     int max_nlayers = 50;
0093     std::vector<std::unordered_map<std::pair<int, int>, GridData, PairHashFunction>> group_hits(max_nlayers);
0094     for (const auto& h : hits) {
0095       auto k = h.getLayer();
0096       if ((int)k >= max_nlayers) {
0097         continue;
0098       }
0099       auto& layer     = group_hits[k];
0100       const auto& pos = h.getPosition();
0101 
0102       // cylindrical r
0103       const float rc   = edm4hep::utils::magnitudeTransverse(pos);
0104       const double eta = edm4hep::utils::eta(pos);
0105       const double phi = edm4hep::utils::angleAzimuthal(pos);
0106 
0107       const auto grid = std::pair<int, int>{pos2grid(eta, m_etaSize), pos2grid(phi, m_phiSize)};
0108       auto it         = layer.find(grid);
0109       // merge energy
0110       if (it != layer.end()) {
0111         auto& data = it->second;
0112         data.nHits += 1;
0113         data.energy += h.getEnergy();
0114         data.energyError += h.getEnergyError() * h.getEnergyError();
0115         data.time += h.getTime();
0116         data.timeError += h.getTimeError() * h.getTimeError();
0117       } else {
0118         layer[grid] = GridData{1,
0119                                rc,
0120                                h.getEnergy(),
0121                                h.getEnergyError() * h.getEnergyError(),
0122                                h.getTime(),
0123                                h.getTimeError() * h.getTimeError(),
0124                                h.getSector()};
0125       }
0126     }
0127 
0128     // convert grid data back to hits
0129     for (const auto& [i, layer] : Jug::Utils::Enumerate(group_hits)) {
0130       for (const auto& [grid, data] : layer) {
0131         const double eta   = grid2pos(grid.first, m_etaSize);
0132         const double phi   = grid2pos(grid.second, m_phiSize);
0133         const double theta = edm4hep::utils::etaToAngle(eta);
0134         const double z     = cotan(theta) * data.rc;
0135         const float r      = std::hypot(data.rc, z);
0136         const auto pos     = edm4hep::utils::sphericalToVector(r, theta, phi);
0137         auto oh            = ohits.create();
0138         oh.setEnergy(data.energy);
0139         oh.setEnergyError(std::sqrt(data.energyError));
0140         oh.setTime(data.time / data.nHits);
0141         oh.setTimeError(std::sqrt(data.timeError));
0142         oh.setPosition(pos);
0143         oh.setLayer(i);
0144         oh.setSector(data.sector);
0145       }
0146     }
0147     return StatusCode::SUCCESS;
0148   }
0149 
0150 private:
0151   static int pos2grid(float pos, float size, float offset = 0.) {
0152     return std::floor((pos - offset) / size);
0153   }
0154   static int grid2pos(float grid, float size, float offset = 0.) {
0155     return (grid + 0.5) * size + offset;
0156   }
0157   static float cotan(float theta) {
0158     if (std::abs(std::sin(theta)) < 1e-6) {
0159       return 0.;
0160     } else {
0161       return 1. / std::tan(theta);
0162     }
0163   }
0164 }; // class ImagingPixelMerger
0165 
0166 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
0167 DECLARE_COMPONENT(ImagingPixelMerger)
0168 
0169 } // namespace Jug::Reco