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 <AthContainers/Decorator.h>
#include <Math/Vector4D.h>
#include <TMath.h>
#include <xAODEventInfo/EventInfo.h>
 
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <vector>
 
namespace CP {

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::unordered_map<std::string, SG::Decorator<float>> floatDecorators;
  std::unordered_map<std::string, SG::Decorator<std::vector<float>>>
      vectorfloatDecorators;
  std::unordered_map<std::string, SG::Decorator<int>> intDecorators;
  std::unordered_map<std::string, SG::Decorator<std::vector<int>>>
      vectorintDecorators;

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

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

  void initializeFloatDecorator(const std::string& name) {
    const std::string key = fullName(name);
    floatDecorators.try_emplace(key, key);
    clearCaches();
  }

  void initializeVectorFloatDecorator(const std::string& name) {
    const std::string key = fullName(name);
    vectorfloatDecorators.try_emplace(key, key);
    clearCaches();
  }

  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 initializeParticleDecorators(const std::string& prefix) {
    initializePtEtaPhiMDecorator(prefix);
    initializeIntDecorator(prefix + "_pdgId");
  }

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

  void initializeVectorParticleDecorators(const std::string& prefix) {
    initializeVectorPtEtaPhiMDecorator(prefix);
    initializeVectorIntDecorator(prefix + "_pdgId");
  }

  void initializeIntDecorator(const std::string& name) {
    const std::string key = fullName(name);
    intDecorators.try_emplace(key, key);
    clearCaches();
  }

  void initializeVectorIntDecorator(const std::string& name) {
    const std::string key = fullName(name);
    vectorintDecorators.try_emplace(key, key);
    clearCaches();
  }

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

  const SG::Decorator<int>* getIntDecorator(const std::string& name) const {
    return findDecorator(intDecorators, name);
  }

  const SG::Decorator<std::vector<int>>* getVectorIntDecorator(
      const std::string& name) const {
    return findDecorator(vectorintDecorators, name);
  }

  const SG::Decorator<float>* getFloatDecorator(const std::string& name) const {
    return findDecorator(floatDecorators, name);
  }

  const SG::Decorator<std::vector<float>>* getVectorFloatDecorator(
      const std::string& name) const {
    return findDecorator(vectorfloatDecorators, name);
  }
 
  // ── Decorate methods ──────────────────────────────────────────────────────
  // All write to m_eventInfo. Call setEventInfo() once per event first.

  void decorateDefault(const std::string& prefix) {
    const ParticleDecorators& decs = particleDecorators(prefix);
    FillDefaultParticleInfo(*decs.pt, *decs.eta, *decs.phi, *decs.m,
                            pdgIdDecorator(decs, prefix), m_eventInfo);
  }

  void decorateVectorDefault(const std::string& prefix) {
    const VectorParticleDecorators& decs = vectorParticleDecorators(prefix);
    FillDefaultVectorParticleInfo(*decs.pt, *decs.eta, *decs.phi, *decs.m,
                                  *decs.pdgId, m_eventInfo);
  }

  void decorateDefaultNoPdgId(const std::string& prefix) {
    const ParticleDecorators& decs = particleDecorators(prefix);
    FillDefaultParticleInfo(*decs.pt, *decs.eta, *decs.phi, *decs.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) {
    const ParticleDecorators& decs = particleDecorators(prefix);
    FillParticleInfo(*decs.pt, *decs.eta, *decs.phi, *decs.m, p, m_eventInfo);
  }

  void decorateParticle(const std::string& prefix,
                        const ROOT::Math::PtEtaPhiMVector& p, int pdgId) {
    const ParticleDecorators& decs = particleDecorators(prefix);
    FillParticleInfo(*decs.pt, *decs.eta, *decs.phi, *decs.m,
                     pdgIdDecorator(decs, prefix), p, pdgId, m_eventInfo);
  }

  void decorateVectorParticle(
      const std::string& prefix,
      const std::vector<ROOT::Math::PtEtaPhiMVector>& vec_p,
      const std::vector<int>& vec_pdgId) {
    const VectorParticleDecorators& decs = vectorParticleDecorators(prefix);
    for (size_t i = 0; i < vec_p.size(); i++) {
      FillVectorParticleInfo(*decs.pt, *decs.eta, *decs.phi, *decs.m,
                             *decs.pdgId, vec_p.at(i), vec_pdgId.at(i),
                             m_eventInfo);
    }
  }
 
 private:
  struct ParticleDecorators {
    const SG::Decorator<float>* pt{nullptr};
    const SG::Decorator<float>* eta{nullptr};
    const SG::Decorator<float>* phi{nullptr};
    const SG::Decorator<float>* m{nullptr};
    const SG::Decorator<int>* pdgId{nullptr};
  };

  struct VectorParticleDecorators {
    const SG::Decorator<std::vector<float>>* pt{nullptr};
    const SG::Decorator<std::vector<float>>* eta{nullptr};
    const SG::Decorator<std::vector<float>>* phi{nullptr};
    const SG::Decorator<std::vector<float>>* m{nullptr};
    const SG::Decorator<std::vector<int>>* pdgId{nullptr};
  };
 
  std::string m_prefix;  
  const xAOD::EventInfo* m_eventInfo{
      nullptr};  

  std::unordered_map<std::string, ParticleDecorators> m_particleCache;
  std::unordered_map<std::string, VectorParticleDecorators>
      m_vectorParticleCache;

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

  template <typename T>
  const T* findDecorator(const std::unordered_map<std::string, T>& decorators,
                         const std::string& name) const {
    const std::string key = fullName(name);
    auto it = decorators.find(key);
    if (it != decorators.end())
      return &it->second;
    throw std::runtime_error("Decorator with name " + key + " not found.");
  }

  void clearCaches() {
    m_particleCache.clear();
    m_vectorParticleCache.clear();
  }

  const ParticleDecorators& particleDecorators(const std::string& prefix) {
    auto it = m_particleCache.find(prefix);
    if (it != m_particleCache.end())
      return it->second;
    ParticleDecorators decs;
    decs.pt = getFloatDecorator(prefix + "_pt");
    decs.eta = getFloatDecorator(prefix + "_eta");
    decs.phi = getFloatDecorator(prefix + "_phi");
    decs.m = getFloatDecorator(prefix + "_m");
    auto pdgIdIt = intDecorators.find(fullName(prefix + "_pdgId"));
    if (pdgIdIt != intDecorators.end())
      decs.pdgId = &pdgIdIt->second;
    return m_particleCache.emplace(prefix, decs).first->second;
  }

  const VectorParticleDecorators& vectorParticleDecorators(
      const std::string& prefix) {
    auto it = m_vectorParticleCache.find(prefix);
    if (it != m_vectorParticleCache.end())
      return it->second;
    VectorParticleDecorators decs;
    decs.pt = getVectorFloatDecorator(prefix + "_pt");
    decs.eta = getVectorFloatDecorator(prefix + "_eta");
    decs.phi = getVectorFloatDecorator(prefix + "_phi");
    decs.m = getVectorFloatDecorator(prefix + "_m");
    decs.pdgId = getVectorIntDecorator(prefix + "_pdgId");
    return m_vectorParticleCache.emplace(prefix, decs).first->second;
  }

  const SG::Decorator<int>& pdgIdDecorator(const ParticleDecorators& decs,
                                           const std::string& prefix) const {
    if (!decs.pdgId)
      throw std::runtime_error("Decorator with name " +
                               fullName(prefix + "_pdgId") + " not found.");
    return *decs.pdgId;
  }
};
 
}  // namespace CP
 
#endif  // DECORATORHELPERS_H