src/components/ImagingPixelMerger.cpp

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