HGCalReco/interface/Trackster.h

0001 // Author: Felice Pantaleo - felice.pantaleo@cern.ch
0002 // Date: 09/2018
0003
0004 #ifndef DataFormats_HGCalReco_Trackster_h
0005 #define DataFormats_HGCalReco_Trackster_h
0006
0007 #include <array>
0008 #include <vector>
0009 #include "DataFormats/Provenance/interface/ProductID.h"
0010 #include "DataFormats/Math/interface/Vector3D.h"
0011
0012 #include <Eigen/Core>
0013
0014 // A Trackster is a Direct Acyclic Graph created when
0015 // pattern recognition algorithms connect hits or
0016 // layer clusters together in a 3D object.
0017
0018 namespace ticl {
0019   class Trackster {
0020   public:
0021     typedef math::XYZVectorF Vector;
0022
0023     enum IterationIndex { TRKEM = 0, EM, TRKHAD, HAD, MIP, SIM, SIM_CP };
0024
0025     // types considered by the particle identification
0026     enum class ParticleType {
0027       photon = 0,
0028       electron,
0029       muon,
0030       neutral_pion,
0031       charged_hadron,
0032       neutral_hadron,
0033       ambiguous,
0034       unknown,
0035     };
0036
0037     enum class PCAOrdering { ascending = 0, descending };
0038
0039     Trackster()
0040         : barycenter_({0.f, 0.f, 0.f}),
0041           regressed_energy_(0.f),
0042           raw_energy_(0.f),
0043           boundTime_(0.f),
0044           time_(0.f),
0045           timeError_(-1.f),
0046           id_probabilities_{},
0047           raw_pt_(0.f),
0048           raw_em_pt_(0.f),
0049           raw_em_energy_(0.f),
0050           seedIndex_(-1),
0051           eigenvalues_{},
0052           sigmas_{},
0053           sigmasPCA_{},
0054           iterationIndex_(0) {}
0055
0056     inline void setIteration(const Trackster::IterationIndex i) { iterationIndex_ = i; }
0057     std::vector<unsigned int> &vertices() { return vertices_; }
0058     std::vector<float> &vertex_multiplicity() { return vertex_multiplicity_; }
0059     std::vector<std::array<unsigned int, 2> > &edges() { return edges_; }
0060     inline void setSeed(edm::ProductID pid, int index) {
0061       seedID_ = pid;
0062       seedIndex_ = index;
0063     }
0064     inline void setTimeAndError(float t, float tError) {
0065       time_ = t;
0066       timeError_ = tError;
0067     }
0068     inline void setRegressedEnergy(float value) { regressed_energy_ = value; }
0069     inline void setRawEnergy(float value) { raw_energy_ = value; }
0070     inline void addToRawEnergy(float value) { raw_energy_ += value; }
0071     inline void setRawEmEnergy(float value) { raw_em_energy_ = value; }
0072     inline void addToRawEmEnergy(float value) { raw_em_energy_ += value; }
0073     inline void setRawPt(float value) { raw_pt_ = value; }
0074     inline void setRawEmPt(float value) { raw_em_pt_ = value; }
0075     inline void setBarycenter(Vector value) { barycenter_ = value; }
0076     inline void setTrackIdx(int index) { track_idx_ = index; }
0077     int trackIdx() const { return track_idx_; }
0078     inline bool isHadronic(float th = 0.5f) const {
0079       return id_probability(Trackster::ParticleType::photon) + id_probability(Trackster::ParticleType::electron) < th;
0080     }
0081     inline void mergeTracksters(const Trackster &other) {
0082       *this += other;
0083
0084       //remove duplicates
0085       removeDuplicates();
0086       zeroProbabilities();
0087     }
0088
0089     inline void mergeTracksters(const std::vector<Trackster> &others) {
0090       for (auto &other : others) {
0091         *this += other;
0092       }
0093
0094       //remove duplicates
0095       removeDuplicates();
0096       zeroProbabilities();
0097     }
0098     inline void fillPCAVariables(Eigen::Vector3f const &eigenvalues,
0099                                  Eigen::Matrix3f const &eigenvectors,
0100                                  Eigen::Vector3f const &sigmas,
0101                                  Eigen::Vector3f const &sigmasEigen,
0102                                  size_t pcadimension,
0103                                  PCAOrdering order) {
0104       int original_index = 0;
0105       for (size_t i = 0; i < pcadimension; ++i) {
0106         sigmas_[i] = std::sqrt(sigmas[i]);
0107         // Reverse the order, since Eigen gives back the eigevalues in
0108         // **increasing order** while we store them in **descreasing order**.
0109         original_index = i;
0110         if (order == PCAOrdering::ascending)
0111           original_index = pcadimension - i - 1;
0112         eigenvalues_[i] = (float)eigenvalues[original_index];
0113         eigenvectors_[i] = ticl::Trackster::Vector(
0114             eigenvectors(0, original_index), eigenvectors(1, original_index), eigenvectors(2, original_index));
0115         sigmasPCA_[i] = std::sqrt(sigmasEigen[original_index]);
0116       }
0117       original_index = 0;
0118       if (order == PCAOrdering::ascending)
0119         original_index = pcadimension - 1;
0120       if (eigenvectors_[0].z() * barycenter_.z() < 0.0) {
0121         eigenvectors_[0] = -ticl::Trackster::Vector(
0122             eigenvectors(0, original_index), eigenvectors(1, original_index), eigenvectors(2, original_index));
0123       }
0124
0125       // Now also update the pt part of the Trackster, using the PCA as direction
0126       raw_pt_ = std::sqrt((eigenvectors_[0].Unit() * raw_energy_).perp2());
0127       raw_em_pt_ = std::sqrt((eigenvectors_[0].Unit() * raw_em_energy_).perp2());
0128     }
0129     void zeroProbabilities() {
0130       for (auto &p : id_probabilities_) {
0131         p = 0.f;
0132       }
0133     }
0134     inline void setProbabilities(float *probs) {
0135       for (float &p : id_probabilities_) {
0136         p = *(probs++);
0137       }
0138     }
0139     inline void setIdProbability(ParticleType type, float value) { id_probabilities_[int(type)] = value; }
0140
0141     inline void setBoundaryTime(float t) { boundTime_ = t; };
0142
0143     inline const Trackster::IterationIndex ticlIteration() const { return (IterationIndex)iterationIndex_; }
0144     inline const std::vector<unsigned int> &vertices() const { return vertices_; }
0145     inline const unsigned int vertices(int index) const { return vertices_[index]; }
0146     inline const std::vector<float> &vertex_multiplicity() const { return vertex_multiplicity_; }
0147     inline const float vertex_multiplicity(int index) const { return vertex_multiplicity_[index]; }
0148     inline const std::vector<std::array<unsigned int, 2> > &edges() const { return edges_; }
0149     inline const edm::ProductID &seedID() const { return seedID_; }
0150     inline const int seedIndex() const { return seedIndex_; }
0151     inline const float time() const { return time_; }
0152     inline const float timeError() const { return timeError_; }
0153     inline const float regressed_energy() const { return regressed_energy_; }
0154     inline const float raw_energy() const { return raw_energy_; }
0155     inline const float raw_em_energy() const { return raw_em_energy_; }
0156     inline const float raw_pt() const { return raw_pt_; }
0157     inline const float raw_em_pt() const { return raw_em_pt_; }
0158     inline const float boundaryTime() const { return boundTime_; };
0159     inline const Vector &barycenter() const { return barycenter_; }
0160     inline const std::array<float, 3> &eigenvalues() const { return eigenvalues_; }
0161     inline const std::array<Vector, 3> &eigenvectors() const { return eigenvectors_; }
0162     inline const Vector &eigenvectors(int index) const { return eigenvectors_[index]; }
0163     inline const std::array<float, 3> &sigmas() const { return sigmas_; }
0164     inline const std::array<float, 3> &sigmasPCA() const { return sigmasPCA_; }
0165     inline const std::array<float, 8> &id_probabilities() const { return id_probabilities_; }
0166     inline const float id_probabilities(int index) const { return id_probabilities_[index]; }
0167
0168     // convenience method to return the ID probability for a certain particle type
0169     inline float id_probability(ParticleType type) const {
0170       // probabilities are stored in the same order as defined in the ParticleType enum
0171       return id_probabilities_[(int)type];
0172     }
0173
0174   private:
0175     Vector barycenter_;
0176     float regressed_energy_;
0177     float raw_energy_;
0178     // -99, -1 if not available. ns units otherwise
0179     float boundTime_;
0180     float time_;
0181     float timeError_;
0182
0183     // trackster ID probabilities
0184     std::array<float, 8> id_probabilities_;
0185
0186     // The vertices of the DAG are the indices of the
0187     // 2d objects in the global collection
0188     std::vector<unsigned int> vertices_;
0189     std::vector<float> vertex_multiplicity_;
0190     float raw_pt_;
0191     float raw_em_pt_;
0192     float raw_em_energy_;
0193
0194     // Product ID of the seeding collection used to create the Trackster.
0195     // For GlobalSeeding the ProductID is set to 0. For track-based seeding
0196     // this is the ProductID of the track-collection used to create the
0197     // seeding-regions.
0198     edm::ProductID seedID_;
0199
0200     // For Global Seeding the index is fixed to one. For track-based seeding,
0201     // the index is the index of the track originating the seeding region that
0202     // created the trackster. For track-based seeding the pointer to the track
0203     // can be cooked using the previous ProductID and this index.
0204     int seedIndex_;
0205     int track_idx_ = -1;
0206
0207     std::array<Vector, 3> eigenvectors_;
0208     std::array<float, 3> eigenvalues_;
0209     std::array<float, 3> sigmas_;
0210     std::array<float, 3> sigmasPCA_;
0211
0212     // The edges connect two vertices together in a directed doublet
0213     // ATTENTION: order matters!
0214     // A doublet generator should create edges in which:
0215     // the first element is on the inner layer and
0216     // the outer element is on the outer layer.
0217     std::vector<std::array<unsigned int, 2> > edges_;
0218
0219     // TICL iteration producing the trackster
0220     uint8_t iterationIndex_;
0221     inline void removeDuplicates() {
0222       auto vtx_sorted{vertices_};
0223       std::sort(std::begin(vtx_sorted), std::end(vtx_sorted));
0224       for (unsigned int iLC = 1; iLC < vtx_sorted.size(); ++iLC) {
0225         if (vtx_sorted[iLC] == vtx_sorted[iLC - 1]) {
0226           // Clean up duplicate LCs
0227           const auto lcIdx = vtx_sorted[iLC];
0228           const auto firstEl = std::find(vertices_.begin(), vertices_.end(), lcIdx);
0229           const auto firstPos = std::distance(std::begin(vertices_), firstEl);
0230           auto iDup = std::find(std::next(firstEl), vertices_.end(), lcIdx);
0231           while (iDup != vertices_.end()) {
0232             vertex_multiplicity_.erase(vertex_multiplicity_.begin() + std::distance(std::begin(vertices_), iDup));
0233             vertices_.erase(iDup);
0234             vertex_multiplicity_[firstPos] -= 1;
0235             iDup = std::find(std::next(firstEl), vertices_.end(), lcIdx);
0236           };
0237         }
0238       }
0239     }
0240     inline void operator+=(const Trackster &other) {
0241       // use getters on other
0242       raw_energy_ += other.raw_energy();
0243       raw_em_energy_ += other.raw_em_energy();
0244       raw_pt_ += other.raw_pt();
0245       raw_em_pt_ += other.raw_em_pt();
0246       // add vertices and multiplicities
0247       std::copy(std::begin(other.vertices()), std::end(other.vertices()), std::back_inserter(vertices_));
0248       std::copy(std::begin(other.vertex_multiplicity()),
0249                 std::end(other.vertex_multiplicity()),
0250                 std::back_inserter(vertex_multiplicity_));
0251     }
0252   };
0253
0254   typedef std::vector<Trackster> TracksterCollection;
0255 }  // namespace ticl
0256 #endif