File indexing completed on 2024-08-21 04:46:19
0001
0002
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
0015
0016
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
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
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
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
0108
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
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
0169 inline float id_probability(ParticleType type) const {
0170
0171 return id_probabilities_[(int)type];
0172 }
0173
0174 private:
0175 Vector barycenter_;
0176 float regressed_energy_;
0177 float raw_energy_;
0178
0179 float boundTime_;
0180 float time_;
0181 float timeError_;
0182
0183
0184 std::array<float, 8> id_probabilities_;
0185
0186
0187
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
0195
0196
0197
0198 edm::ProductID seedID_;
0199
0200
0201
0202
0203
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
0213
0214
0215
0216
0217 std::vector<std::array<unsigned int, 2> > edges_;
0218
0219
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
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
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
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 }
0256 #endif