Charge field (#132)

* add charge field as default for track and particle

* fix bug of prev. commit

* check if field to be added doesn't already exist;
bugfix: add setter for charge field

* apply clang format

Author: lubynets

core/BranchConfig.cpp

@@ -25,6 +25,7 @@ BranchConfig::BranchConfig(std::string name, DetType type) : name_(std::move(nam
     VectorConfig<float>::AddField("phi", TrackFields::kPhi, "azimuthal angle");
     VectorConfig<float>::AddField("eta", TrackFields::kEta, "pseudorapidity");
     VectorConfig<float>::AddField("p", TrackFields::kP, "GeV/c");
+    VectorConfig<int>::AddField("q", TrackFields::kQ, "Charge of the track (or its sign when absolute value unknown)");
     VectorConfig<int>::AddField("id", TrackFields::kId, "unique id");
   } else if (type_ == DetType::kParticle) {
     VectorConfig<float>::AddField("px", ParticleFields::kPx, "GeV/c");
@@ -38,6 +39,7 @@ BranchConfig::BranchConfig(std::string name, DetType type) : name_(std::move(nam
     VectorConfig<float>::AddField("E", ParticleFields::kEnergy, "full energy, GeV");
     VectorConfig<float>::AddField("T", ParticleFields::kKineticEnergy, "kinetic energy, GeV");
     VectorConfig<float>::AddField("rapidity", ParticleFields::kRapidity, "in Lab. frame");
+    VectorConfig<int>::AddField("q", ParticleFields::kQ, "Charge of the particle");
     VectorConfig<int>::AddField("pid", ParticleFields::kPid, "PDG code");
     VectorConfig<int>::AddField("id", ParticleFields::kId, "unique id");
   } else if (type_ == DetType::kHit) {
@@ -148,6 +150,12 @@ std::vector<std::string> VectorConfig<T>::SplitString(const std::string& input)
   return result;
 }
 
+void BranchConfig::GuaranteeFieldNameVacancy(const std::string& name) const {
+  if (HasField(name)) {
+    throw std::runtime_error("BranchConfig::GuaranteeFieldNameVacancy(): field " + name + " already exists");
+  }
+}
+
 void BranchConfig::Print() const {
   std::cout << "Branch " << name_ << " (id=" << id_ << ") consists of:" << std::endl;
   std::cout << "Floating fields:" << std::endl;

core/BranchConfig.hpp

@@ -134,14 +134,19 @@ class BranchConfig : public VectorConfig<int>, public VectorConfig<float>, publi
   // Setters
   template<typename T>
   void AddField(const std::string& name, const std::string& title = "") {
+    GuaranteeFieldNameVacancy(name);
     VectorConfig<T>::AddField(name, title);
   }
   template<typename T>
   void AddFields(const std::vector<std::string>& names, const std::string& title = "") {
+    for (auto& n : names) {
+      GuaranteeFieldNameVacancy(n);
+    }
     VectorConfig<T>::AddFields(names, title);
   }
   template<typename T>
   void AddField(const std::string& name, ShortInt_t id, const std::string& title = "") {
+    GuaranteeFieldNameVacancy(name);
     VectorConfig<T>::AddField(name, id, title);
   }
 
@@ -176,6 +181,8 @@ class BranchConfig : public VectorConfig<int>, public VectorConfig<float>, publi
  protected:
   void GenerateId();
 
+  void GuaranteeFieldNameVacancy(const std::string& name) const;
+
   std::string name_;
   size_t id_{0};
   DetType type_{DetType(UndefValueShort)};

core/Constants.hpp

@@ -67,7 +67,8 @@ enum TrackFields : ShortInt_t {
   kPy = -5,
   kPz = -6,
   kP = -7,
-  kId = -8
+  kQ = -8,
+  kId = -9
 };
 }
 
@@ -85,7 +86,8 @@ enum ParticleFields : ShortInt_t {
   kP = -10,
   kEnergy = -11,
   kKineticEnergy = -12,
-  kId = -13
+  kQ = -13,
+  kId = -14
 };
 }
 

core/Particle.cpp

@@ -9,5 +9,9 @@ void Particle::SetPid(PdgCode_t pid) {
   pid_ = pid;
   if (mass_ == -1000.f)
     mass_ = GetMassByPdgId(pid);
+
+  if (charge_ == -1000) {
+    charge_ = GetChargeByPdgId(pid);
+  }
 }
 }// namespace AnalysisTree

core/Particle.hpp

@@ -55,6 +55,7 @@ class Particle : public Track {
         case ParticleFields::kPx: return GetPx();
         case ParticleFields::kPy: return GetPy();
         case ParticleFields::kPz: return GetPz();
+        case ParticleFields::kQ: return GetCharge();
         case ParticleFields::kId: return GetId();
         default: throw std::out_of_range("Particle::GetField - Index " + std::to_string(iField) + " is not found");
       }
@@ -71,6 +72,7 @@ class Particle : public Track {
         case ParticleFields::kPy: py_ = value; break;
         case ParticleFields::kPz: pz_ = value; break;
         case ParticleFields::kMass: mass_ = value; break;
+        case ParticleFields::kQ: charge_ = value; break;
         case ParticleFields::kPid: SetPid(value); break;
         case ParticleFields::kId: break;
         case ParticleFields::kP: /*throw std::runtime_error("Cannot set transient fields");*/ break;

core/Track.cpp

@@ -41,4 +41,20 @@ float Track::GetMassByPdgId(PdgCode_t pdg) {
   }
 }
 
+int Track::GetChargeByPdgId(PdgCode_t pdg) {
+
+  if (pdg > 1000000000) {//100ZZZAAA0
+    auto Z = (pdg % 10000000) / 10000;
+    return Z;
+  }
+  auto db = TDatabasePDG::Instance();
+  auto particle = db->GetParticle(pdg);
+
+  if (particle) {
+    return int(particle->Charge() / 3);
+  } else {
+    throw std::runtime_error("Mass of " + std::to_string(pdg) + " is not known");
+  }
+}
+
 }// namespace AnalysisTree

core/Track.hpp

@@ -63,13 +63,16 @@ class Track : public Container {
     pz_ = momentum.Pz();
   }
 
+  void SetCharge(Int_t charge) noexcept { charge_ = charge; }
+
   ANALYSISTREE_ATTR_NODISCARD inline Floating_t GetPx() const noexcept { return px_; }
   ANALYSISTREE_ATTR_NODISCARD inline Floating_t GetPy() const noexcept { return py_; }
   ANALYSISTREE_ATTR_NODISCARD inline Floating_t GetPz() const noexcept { return pz_; }
   ANALYSISTREE_ATTR_NODISCARD inline Floating_t GetPt() const noexcept { return sqrt(px_ * px_ + py_ * py_); }
   ANALYSISTREE_ATTR_NODISCARD inline Floating_t GetPhi() const noexcept { return atan2(py_, px_); }
   ANALYSISTREE_ATTR_NODISCARD inline Floating_t GetEta() const noexcept { return 0.5f * log((GetP() + pz_) / (GetP() - pz_)); }
   ANALYSISTREE_ATTR_NODISCARD inline Floating_t GetP() const noexcept { return sqrt(px_ * px_ + py_ * py_ + pz_ * pz_); }
+  ANALYSISTREE_ATTR_NODISCARD inline Integer_t GetCharge() const noexcept { return charge_; }
 
   /**
   * @return 3d-momentum of a track
@@ -120,6 +123,7 @@ class Track : public Container {
         case TrackFields::kPx: return GetPx();
         case TrackFields::kPy: return GetPy();
         case TrackFields::kPz: return GetPz();
+        case TrackFields::kQ: return GetCharge();
         case TrackFields::kId: return GetId();
         default: throw std::out_of_range("Track::GetField - Index " + std::to_string(id) + " is not found");
       }
@@ -135,6 +139,7 @@ class Track : public Container {
         case TrackFields::kPx: px_ = value; break;
         case TrackFields::kPy: py_ = value; break;
         case TrackFields::kPz: pz_ = value; break;
+        case TrackFields::kQ: charge_ = value; break;
         case TrackFields::kId: break;
         case TrackFields::kP: /*throw std::runtime_error("Cannot set transient fields");*/ break;
         case TrackFields::kPt: /*throw std::runtime_error("Cannot set transient fields");*/ break;
@@ -152,10 +157,12 @@ class Track : public Container {
 
  protected:
   static float GetMassByPdgId(PdgCode_t pdg);
+  static int GetChargeByPdgId(PdgCode_t pdg);
 
   Floating_t px_{UndefValueFloat};///< x-component of track's momentum
   Floating_t py_{UndefValueFloat};///< y-component of track's momentum
   Floating_t pz_{UndefValueFloat};///< z-component of track's momentum
+  Integer_t charge_{-1000};
 
   ClassDefOverride(Track, 2);
 };

infra/TaskManager.cpp

@@ -150,7 +150,7 @@ void TaskManager::Finish() {
     out_tree_->Write();
     configuration_->Write("Configuration");
     data_header_->Write("DataHeader");
-    if(is_write_hash_info_) WriteCommitInfo();
+    if (is_write_hash_info_) WriteCommitInfo();
     out_file_->Close();
     out_tree_ = nullptr;
     delete out_file_;

infra/TaskManager.hpp

@@ -158,7 +158,7 @@ class TaskManager {
   void SetWriteMode(eBranchWriteMode mode) { write_mode_ = mode; }
   void SetBranchesExclude(std::vector<std::string> brex) { branches_exclude_ = std::move(brex); }
   void SetVerbosityPeriod(int value) { verbosity_period_ = value; }
-  void SetIsWriteHashInfo(bool is=true) { is_write_hash_info_ = is; }
+  void SetIsWriteHashInfo(bool is = true) { is_write_hash_info_ = is; }
   void SetIsUpdateEntryInExec(bool is = true) { is_update_entry_in_exec_ = is; }
 
   void ClearTasks() { tasks_.clear(); }