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