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