DecoratorHelpers.h

Go to the documentation of this file.

/*
  Copyright (C) 2002-2026 CERN for the benefit of the ATLAS collaboration
*/

 
#ifndef DECORATORHELPERS_H
#define DECORATORHELPERS_H
 
#include <Math/Vector4D.h>
#include <TMath.h>
#include <xAODEventInfo/EventInfo.h>

template <typename T, typename U, typename V>
void FillParticleInfo(const SG::Decorator<T>& dec, const U& value,
                      const V* object) {
  dec(*object) = value;
}

template <typename T, typename U, typename V>
void FillVectorParticleInfo(const SG::Decorator<std::vector<T>>& dec,
                            const U& value, const V* object) {
  dec(*object).push_back(value);
}

template <typename T, typename U>
void FillParticleInfo(const SG::Decorator<T>& dec_pt,
                      const SG::Decorator<T>& dec_eta,
                      const SG::Decorator<T>& dec_phi,
                      const SG::Decorator<T>& dec_m,
                      const ROOT::Math::PtEtaPhiMVector& particle, U* object) {
  FillParticleInfo(dec_pt, particle.Pt(), object);
  FillParticleInfo(dec_eta, particle.Eta(), object);
  FillParticleInfo(dec_phi, particle.Phi(), object);
  FillParticleInfo(dec_m, particle.M(), object);
}

template <typename T, typename U>
void FillVectorParticleInfo(const SG::Decorator<std::vector<T>>& dec_pt,
                            const SG::Decorator<std::vector<T>>& dec_eta,
                            const SG::Decorator<std::vector<T>>& dec_phi,
                            const SG::Decorator<std::vector<T>>& dec_m,
                            const ROOT::Math::PtEtaPhiMVector& particle,
                            U* object) {
  FillVectorParticleInfo(dec_pt, particle.Pt(), object);
  FillVectorParticleInfo(dec_eta, particle.Eta(), object);
  FillVectorParticleInfo(dec_phi, particle.Phi(), object);
  FillVectorParticleInfo(dec_m, particle.M(), object);
}

template <typename T, typename U>
void FillParticleInfo(const SG::Decorator<T>& dec_pt,
                      const SG::Decorator<T>& dec_eta,
                      const SG::Decorator<T>& dec_phi,
                      const SG::Decorator<T>& dec_m,
                      const SG::Decorator<int>& dec_pdgId,
                      const ROOT::Math::PtEtaPhiMVector& particle, int pdgId,
                      U* object) {
  FillParticleInfo(dec_pdgId, pdgId, object);
  FillParticleInfo(dec_pt, dec_eta, dec_phi, dec_m, particle, object);
}

template <typename T, typename U>
void FillVectorParticleInfo(const SG::Decorator<std::vector<T>>& dec_pt,
                            const SG::Decorator<std::vector<T>>& dec_eta,
                            const SG::Decorator<std::vector<T>>& dec_phi,
                            const SG::Decorator<std::vector<T>>& dec_m,
                            const SG::Decorator<std::vector<int>>& dec_pdgId,
                            const ROOT::Math::PtEtaPhiMVector& particle,
                            int pdgId, U* object) {
  FillVectorParticleInfo(dec_pdgId, pdgId, object);
  FillVectorParticleInfo(dec_pt, dec_eta, dec_phi, dec_m, particle, object);
}

template <typename T, typename U>
void FillDefaultParticleInfo(const SG::Decorator<T>& dec_pt,
                             const SG::Decorator<T>& dec_eta,
                             const SG::Decorator<T>& dec_phi,
                             const SG::Decorator<T>& dec_m, U* object) {
  FillParticleInfo(dec_pt, -1, object);
  FillParticleInfo(dec_eta, -999, object);
  FillParticleInfo(dec_phi, -999, object);
  FillParticleInfo(dec_m, -1, object);
}

template <typename T, typename U>
void FillDefaultVectorParticleInfo(const SG::Decorator<std::vector<T>>& dec_pt,
                                   const SG::Decorator<std::vector<T>>& dec_eta,
                                   const SG::Decorator<std::vector<T>>& dec_phi,
                                   const SG::Decorator<std::vector<T>>& dec_m,
                                   U* object) {
  dec_pt(*object) = {};
  dec_eta(*object) = {};
  dec_phi(*object) = {};
  dec_m(*object) = {};
}

template <typename T, typename U>
void FillDefaultParticleInfo(const SG::Decorator<T>& dec_pt,
                             const SG::Decorator<T>& dec_eta,
                             const SG::Decorator<T>& dec_phi,
                             const SG::Decorator<T>& dec_m,
                             const SG::Decorator<int>& dec_pdgId, U* object) {
  FillParticleInfo(dec_pdgId, 0, object);
  FillDefaultParticleInfo(dec_pt, dec_eta, dec_phi, dec_m, object);
}

template <typename T, typename U>
void FillDefaultVectorParticleInfo(
    const SG::Decorator<std::vector<T>>& dec_pt,
    const SG::Decorator<std::vector<T>>& dec_eta,
    const SG::Decorator<std::vector<T>>& dec_phi,
    const SG::Decorator<std::vector<T>>& dec_m,
    const SG::Decorator<std::vector<int>>& dec_pdgId, U* object) {
  dec_pdgId(*object) = {};
  FillDefaultVectorParticleInfo(dec_pt, dec_eta, dec_phi, dec_m, object);
}

struct PartonDecorator {
  std::map<std::string, std::unique_ptr<SG::Decorator<float>>> floatDecorators;
  std::map<std::string, std::unique_ptr<SG::Decorator<std::vector<float>>>>
      vectorfloatDecorators;
  std::map<std::string, std::unique_ptr<SG::Decorator<int>>> intDecorators;
  std::map<std::string, std::unique_ptr<SG::Decorator<std::vector<int>>>>
      vectorintDecorators;

  void setPrefix(const std::string& prefix) { m_prefix = prefix; }

  void setEventInfo(const xAOD::EventInfo* eventInfo) {
    m_eventInfo = eventInfo;
  }
 
  // ── Initialize methods ────────────────────────────────────────────────────

  void initializeFloatDecorator(const std::string& name) {
    const std::string key = fullName(name);
    floatDecorators[key] = std::make_unique<SG::Decorator<float>>(key);
  }

  void initializeVectorFloatDecorator(const std::string& name) {
    const std::string key = fullName(name);
    vectorfloatDecorators[key] =
        std::make_unique<SG::Decorator<std::vector<float>>>(key);
  }

  void initializeFloatDecorator(const std::vector<std::string>& names) {
    for (const auto& name : names)
      initializeFloatDecorator(name);
  }

  void initializePtEtaPhiMDecorator(const std::string& prefix) {
    initializeFloatDecorator(prefix + "_pt");
    initializeFloatDecorator(prefix + "_eta");
    initializeFloatDecorator(prefix + "_phi");
    initializeFloatDecorator(prefix + "_m");
  }

  void initializeVectorPtEtaPhiMDecorator(const std::string& prefix) {
    initializeVectorFloatDecorator(prefix + "_pt");
    initializeVectorFloatDecorator(prefix + "_eta");
    initializeVectorFloatDecorator(prefix + "_phi");
    initializeVectorFloatDecorator(prefix + "_m");
  }

  void initializeIntDecorator(const std::string& name) {
    const std::string key = fullName(name);
    intDecorators[key] = std::make_unique<SG::Decorator<int>>(key);
  }

  void initializeVectorIntDecorator(const std::string& name) {
    const std::string key = fullName(name);
    vectorintDecorators[key] =
        std::make_unique<SG::Decorator<std::vector<int>>>(key);
  }

  void initializeIntDecorator(const std::vector<std::string>& names) {
    for (const auto& name : names)
      initializeIntDecorator(name);
  }
 
  // ── Getter methods ────────────────────────────────────────────────────────

  SG::Decorator<int>* getIntDecorator(const std::string& name) const {
    const std::string key = fullName(name);
    auto it = intDecorators.find(key);
    if (it != intDecorators.end())
      return it->second.get();
    throw std::runtime_error("Decorator with name " + key + " not found.");
  }

  SG::Decorator<std::vector<int>>* getVectorIntDecorator(
      const std::string& name) const {
    const std::string key = fullName(name);
    auto it = vectorintDecorators.find(key);
    if (it != vectorintDecorators.end())
      return it->second.get();
    throw std::runtime_error("Decorator with name " + key + " not found.");
  }

  SG::Decorator<float>* getFloatDecorator(const std::string& name) const {
    const std::string key = fullName(name);
    auto it = floatDecorators.find(key);
    if (it != floatDecorators.end())
      return it->second.get();
    throw std::runtime_error("Decorator with name " + key + " not found.");
  }

  SG::Decorator<std::vector<float>>* getVectorFloatDecorator(
      const std::string& name) const {
    const std::string key = fullName(name);
    auto it = vectorfloatDecorators.find(key);
    if (it != vectorfloatDecorators.end())
      return it->second.get();
    throw std::runtime_error("Decorator with name " + key + " not found.");
  }
 
  // ── Decorate methods ──────────────────────────────────────────────────────
  // All write to m_eventInfo. Call setEventInfo() once per event first.

  void decorateDefault(const std::string& prefix) {
    FillDefaultParticleInfo(
        *getFloatDecorator(prefix + "_pt"), *getFloatDecorator(prefix + "_eta"),
        *getFloatDecorator(prefix + "_phi"), *getFloatDecorator(prefix + "_m"),
        *getIntDecorator(prefix + "_pdgId"), m_eventInfo);
  }

  void decorateVectorDefault(const std::string& prefix) {
    FillDefaultVectorParticleInfo(*getVectorFloatDecorator(prefix + "_pt"),
                                  *getVectorFloatDecorator(prefix + "_eta"),
                                  *getVectorFloatDecorator(prefix + "_phi"),
                                  *getVectorFloatDecorator(prefix + "_m"),
                                  *getVectorIntDecorator(prefix + "_pdgId"),
                                  m_eventInfo);
  }

  void decorateDefaultNoPdgId(const std::string& prefix) {
    FillDefaultParticleInfo(*getFloatDecorator(prefix + "_pt"),
                            *getFloatDecorator(prefix + "_eta"),
                            *getFloatDecorator(prefix + "_phi"),
                            *getFloatDecorator(prefix + "_m"), m_eventInfo);
  }

  void decorateCustom(const std::string& name, float value) {
    FillParticleInfo(*getFloatDecorator(name), value, m_eventInfo);
  }

  void decorateCustom(const std::string& name, int value) {
    FillParticleInfo(*getIntDecorator(name), value, m_eventInfo);
  }

  void decorateParticle(const std::string& prefix,
                        const ROOT::Math::PtEtaPhiMVector& p) {
    FillParticleInfo(*getFloatDecorator(prefix + "_pt"),
                     *getFloatDecorator(prefix + "_eta"),
                     *getFloatDecorator(prefix + "_phi"),
                     *getFloatDecorator(prefix + "_m"), p, m_eventInfo);
  }

  void decorateParticle(const std::string& prefix,
                        const ROOT::Math::PtEtaPhiMVector& p, int pdgId) {
    FillParticleInfo(
        *getFloatDecorator(prefix + "_pt"), *getFloatDecorator(prefix + "_eta"),
        *getFloatDecorator(prefix + "_phi"), *getFloatDecorator(prefix + "_m"),
        *getIntDecorator(prefix + "_pdgId"), p, pdgId, m_eventInfo);
  }

  void decorateVectorParticle(
      const std::string& prefix,
      const std::vector<ROOT::Math::PtEtaPhiMVector>& vec_p,
      const std::vector<int>& vec_pdgId) {
    for (size_t i = 0; i < vec_p.size(); i++) {
      FillVectorParticleInfo(*getVectorFloatDecorator(prefix + "_pt"),
                             *getVectorFloatDecorator(prefix + "_eta"),
                             *getVectorFloatDecorator(prefix + "_phi"),
                             *getVectorFloatDecorator(prefix + "_m"),
                             *getVectorIntDecorator(prefix + "_pdgId"),
                             vec_p.at(i), vec_pdgId.at(i), m_eventInfo);
    }
  }
 
 private:
  std::string m_prefix;  
  const xAOD::EventInfo* m_eventInfo{
      nullptr};  

  std::string fullName(const std::string& name) const {
    return m_prefix.empty() ? name : m_prefix + "_" + name;
  }
};
 
#endif  // DECORATORHELPERS_H