File indexing completed on 2023-03-17 10:50:06
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::XYZVector 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 : iterationIndex_(0),
0041 seedIndex_(-1),
0042 time_(0.f),
0043 timeError_(-1.f),
0044 regressed_energy_(0.f),
0045 raw_energy_(0.f),
0046 raw_em_energy_(0.f),
0047 raw_pt_(0.f),
0048 raw_em_pt_(0.f),
0049 barycenter_({0., 0., 0.}),
0050 eigenvalues_{{0.f, 0.f, 0.f}},
0051 sigmas_{{0.f, 0.f, 0.f}},
0052 sigmasPCA_{{0.f, 0.f, 0.f}} {
0053 zeroProbabilities();
0054 }
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 fillPCAVariables(Eigen::Vector3d &eigenvalues,
0077 Eigen::Matrix3d &eigenvectors,
0078 Eigen::Vector3d &sigmas,
0079 Eigen::Vector3d &sigmasEigen,
0080 size_t pcadimension,
0081 PCAOrdering order) {
0082 int original_index = 0;
0083 for (size_t i = 0; i < pcadimension; ++i) {
0084 sigmas_[i] = std::sqrt(sigmas[i]);
0085
0086
0087 original_index = i;
0088 if (order == PCAOrdering::ascending)
0089 original_index = pcadimension - i - 1;
0090 eigenvalues_[i] = (float)eigenvalues[original_index];
0091 eigenvectors_[i] = ticl::Trackster::Vector(
0092 eigenvectors(0, original_index), eigenvectors(1, original_index), eigenvectors(2, original_index));
0093 sigmasPCA_[i] = std::sqrt(sigmasEigen[original_index]);
0094 }
0095 original_index = 0;
0096 if (order == PCAOrdering::ascending)
0097 original_index = pcadimension - 1;
0098 if (eigenvectors_[0].z() * barycenter_.z() < 0.0) {
0099 eigenvectors_[0] = -ticl::Trackster::Vector(
0100 eigenvectors(0, original_index), eigenvectors(1, original_index), eigenvectors(2, original_index));
0101 }
0102
0103
0104 raw_pt_ = std::sqrt((eigenvectors_[0].Unit() * raw_energy_).perp2());
0105 raw_em_pt_ = std::sqrt((eigenvectors_[0].Unit() * raw_em_energy_).perp2());
0106 }
0107 void zeroProbabilities() {
0108 for (auto &p : id_probabilities_) {
0109 p = 0.f;
0110 }
0111 }
0112 inline void setProbabilities(float *probs) {
0113 for (float &p : id_probabilities_) {
0114 p = *(probs++);
0115 }
0116 }
0117 inline void setIdProbability(ParticleType type, float value) { id_probabilities_[int(type)] = value; }
0118
0119 inline const Trackster::IterationIndex ticlIteration() const { return (IterationIndex)iterationIndex_; }
0120 inline const std::vector<unsigned int> &vertices() const { return vertices_; }
0121 inline const unsigned int vertices(int index) const { return vertices_[index]; }
0122 inline const std::vector<float> &vertex_multiplicity() const { return vertex_multiplicity_; }
0123 inline const float vertex_multiplicity(int index) const { return vertex_multiplicity_[index]; }
0124 inline const std::vector<std::array<unsigned int, 2> > &edges() const { return edges_; }
0125 inline const edm::ProductID &seedID() const { return seedID_; }
0126 inline const int seedIndex() const { return seedIndex_; }
0127 inline const float time() const { return time_; }
0128 inline const float timeError() const { return timeError_; }
0129 inline const float regressed_energy() const { return regressed_energy_; }
0130 inline const float raw_energy() const { return raw_energy_; }
0131 inline const float raw_em_energy() const { return raw_em_energy_; }
0132 inline const float raw_pt() const { return raw_pt_; }
0133 inline const float raw_em_pt() const { return raw_em_pt_; }
0134 inline const Vector &barycenter() const { return barycenter_; }
0135 inline const std::array<float, 3> &eigenvalues() const { return eigenvalues_; }
0136 inline const std::array<Vector, 3> &eigenvectors() const { return eigenvectors_; }
0137 inline const Vector &eigenvectors(int index) const { return eigenvectors_[index]; }
0138 inline const std::array<float, 3> &sigmas() const { return sigmas_; }
0139 inline const std::array<float, 3> &sigmasPCA() const { return sigmasPCA_; }
0140 inline const std::array<float, 8> &id_probabilities() const { return id_probabilities_; }
0141 inline const float id_probabilities(int index) const { return id_probabilities_[index]; }
0142
0143
0144 inline float id_probability(ParticleType type) const {
0145
0146 return id_probabilities_[(int)type];
0147 }
0148
0149 private:
0150
0151 uint8_t iterationIndex_;
0152
0153
0154
0155 std::vector<unsigned int> vertices_;
0156 std::vector<float> vertex_multiplicity_;
0157
0158
0159
0160
0161
0162
0163 std::vector<std::array<unsigned int, 2> > edges_;
0164
0165
0166
0167
0168
0169 edm::ProductID seedID_;
0170
0171
0172
0173
0174
0175 int seedIndex_;
0176
0177
0178
0179
0180
0181
0182 float time_;
0183 float timeError_;
0184
0185
0186 float regressed_energy_;
0187 float raw_energy_;
0188 float raw_em_energy_;
0189 float raw_pt_;
0190 float raw_em_pt_;
0191
0192
0193 Vector barycenter_;
0194 std::array<float, 3> eigenvalues_;
0195 std::array<Vector, 3> eigenvectors_;
0196 std::array<float, 3> sigmas_;
0197 std::array<float, 3> sigmasPCA_;
0198
0199
0200 std::array<float, 8> id_probabilities_;
0201 };
0202
0203 typedef std::vector<Trackster> TracksterCollection;
0204 }
0205 #endif