src/components/CalorimeterHitsEtaPhiProjector.cpp

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2022 Jihee Kim, Sylvester Joosten, Chao Peng, Whitney Armstrong, Wouter Deconinck, Chao Peng
0003
0004 /*
0005  *  A hits converter to prepare dataset for machine learning
0006  *  It converts hits with (x, y, z, E) to (E, eta, phi) layer by layer
0007  *  With a defined grid size and ranges, it merge the hits within one grid and drop-off hits
0008  * out-of-range The capacity of each layer is fixed (padding with zeros), and the hits with least
0009  * energies that exceed the capacity will be discarded.
0010  *
0011  *  Author: Chao Peng (ANL), Jihee Kim, 07/16/2021
0012  *
0013  */
0014 #include <algorithm>
0015 #include <bitset>
0016 #include <unordered_map>
0017
0018 #include "Gaudi/Property.h"
0019 #include "GaudiAlg/GaudiAlgorithm.h"
0020 #include "GaudiAlg/GaudiTool.h"
0021 #include "GaudiAlg/Transformer.h"
0022 #include "GaudiKernel/PhysicalConstants.h"
0023 #include "GaudiKernel/RndmGenerators.h"
0024 #include "GaudiKernel/ToolHandle.h"
0025
0026 #include "DDRec/CellIDPositionConverter.h"
0027 #include "DDRec/Surface.h"
0028 #include "DDRec/SurfaceManager.h"
0029
0030 #include <k4FWCore/DataHandle.h>
0031 #include <k4Interface/IGeoSvc.h>
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 using Point3D = ROOT::Math::XYZPoint;
0039
0040 struct pair_hash {
0041   template <class T1, class T2> std::size_t operator()(const std::pair<T1, T2>& pair) const {
0042     return std::hash<T1>()(pair.first) ^ std::hash<T2>()(pair.second);
0043   }
0044 };
0045
0046 namespace Jug::Reco {
0047
0048 /** Calorimeter eta-phi projector
0049  *
0050  *  A hits converter to prepare dataset for machine learning
0051  *  It converts hits with (x, y, z, E) to (E, eta, phi) layer by layer
0052  *  With a defined grid size and ranges, it merge the hits within one grid and drop-off hits
0053  * out-of-range The capacity of each layer is fixed (padding with zeros), and the hits with least
0054  * energies that exceed the capacity will be discarded.
0055  *
0056  *
0057  * \ingroup reco
0058  */
0059 class CalorimeterHitsEtaPhiProjector : public GaudiAlgorithm {
0060 private:
0061   Gaudi::Property<std::vector<double>> u_gridSizes{this, "gridSizes", {0.001, 0.001 * rad}};
0062   DataHandle<edm4eic::CalorimeterHitCollection> m_inputHitCollection{"inputHitCollection", Gaudi::DataHandle::Reader,
0063                                                                   this};
0064   DataHandle<edm4eic::CalorimeterHitCollection> m_outputHitCollection{"outputHitCollection", Gaudi::DataHandle::Writer,
0065                                                                    this};
0066
0067   double gridSizes[2]{0.0, 0.0};
0068
0069 public:
0070   CalorimeterHitsEtaPhiProjector(const std::string& name, ISvcLocator* svcLoc) : GaudiAlgorithm(name, svcLoc) {
0071     declareProperty("inputHitCollection", m_inputHitCollection, "");
0072     declareProperty("outputHitCollection", m_outputHitCollection, "");
0073   }
0074
0075   StatusCode initialize() override {
0076     if (GaudiAlgorithm::initialize().isFailure()) {
0077       return StatusCode::FAILURE;
0078     }
0079
0080     if (u_gridSizes.size() != 2) {
0081       error() << "Expected 2 values for gridSizes, received " << u_gridSizes.size() << endmsg;
0082       return StatusCode::FAILURE;
0083     }
0084     gridSizes[0] = u_gridSizes.value()[0];
0085     gridSizes[1] = u_gridSizes.value()[1] / rad;
0086
0087     return StatusCode::SUCCESS;
0088   }
0089
0090   StatusCode execute() override {
0091     // Create output collections
0092     auto& mhits = *m_outputHitCollection.createAndPut();
0093
0094     // container
0095     std::unordered_map<std::pair<int64_t, int64_t>, std::vector<edm4eic::CalorimeterHit>, pair_hash> merged_hits;
0096
0097     for (const auto h : *m_inputHitCollection.get()) {
0098       auto bins =
0099           std::make_pair(static_cast<int64_t>(pos2bin(edm4hep::utils::eta(h.getPosition()), gridSizes[0], 0.)),
0100                          static_cast<int64_t>(pos2bin(edm4hep::utils::angleAzimuthal(h.getPosition()), gridSizes[1], 0.)));
0101       merged_hits[bins].push_back(h);
0102     }
0103
0104     for (const auto& [bins, hits] : merged_hits) {
0105       const auto ref = hits.front();
0106       edm4eic::MutableCalorimeterHit hit;
0107       hit.setCellID(ref.getCellID());
0108       // TODO, we can do timing cut to reject noises
0109       hit.setTime(ref.getTime());
0110       double r   = edm4hep::utils::magnitude(ref.getPosition());
0111       double eta = bin2pos(bins.first, gridSizes[0], 0.);
0112       double phi = bin2pos(bins.second, gridSizes[1], 1.);
0113       hit.setPosition(edm4hep::utils::sphericalToVector(r, edm4hep::utils::etaToAngle(eta), phi));
0114       hit.setDimension({static_cast<float>(gridSizes[0]), static_cast<float>(gridSizes[1]), 0.});
0115       // merge energy
0116       hit.setEnergy(0.);
0117       for (const auto& h : hits) {
0118         hit.setEnergy(hit.getEnergy() + h.getEnergy());
0119       }
0120       mhits.push_back(hit);
0121     }
0122
0123     return StatusCode::SUCCESS;
0124   }
0125
0126   static int64_t pos2bin(double val, double cell, double offset = 0.) {
0127     return int64_t(std::floor((val + 0.5 * cell - offset) / cell));
0128   }
0129
0130   static double bin2pos(int64_t bin, double cell, double offset = 0.) { return bin * cell + offset; }
0131
0132 }; // class CalorimeterHitsEtaPhiProjector
0133
0134 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
0135 DECLARE_COMPONENT(CalorimeterHitsEtaPhiProjector)
0136
0137 } // namespace Jug::Reco