MomentumAccessors.h

Go to the documentation of this file.

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

 
 
#ifndef COLUMNAR_CORE_MOMENTUM_ACCESSORS_H
#define COLUMNAR_CORE_MOMENTUM_ACCESSORS_H
 
#include <ColumnarCore/ColumnAccessor.h>
#include <TruthUtils/ParticleConstants.h>
#include <xAODBase/ObjectType.h>
#include <Math/Vector4D.h>
 
namespace columnar
{
 
  namespace Detail
  {
 
    template<ContainerIdConcept CIVal, typename CMVal = ColumnarModeDefault>
    class MomentumAccessorsIParticle
    {
    public:
      using CI = CIVal;
      using CM = CMVal;
 
      double pt (ObjectId<CI,CM> object) const {
        return object.getXAODObject().pt();}
      double eta (ObjectId<CI,CM> object) const {
        return object.getXAODObject().eta();}
      double phi (ObjectId<CI,CM> object) const {
        return object.getXAODObject().phi();}
      double m (ObjectId<CI,CM> object) const {
        return object.getXAODObject().m();}
      double e (ObjectId<CI,CM> object) const {
        return object.getXAODObject().e();}
      double rapidity (ObjectId<CI,CM> object) const {
        return object.getXAODObject().rapidity();}
    };

 
    template<typename CoreAccessors>
    class FullMomentumAccessorsPtEtaPhiM : public CoreAccessors
    {
    public:
      using CI = CoreAccessors::CI;
      using CM = typename CoreAccessors::CM;
 
      using CoreAccessors::CoreAccessors;
 
      double e (ObjectId<CI,CM> object) const
      {
        const double myPt = this->pt (object);
        const double myEta = this->eta (object);
        const double myM = this->m (object);
 
        // not reading phi, not part of the energy calculation
        const ROOT::Math::LorentzVector<ROOT::Math::PtEtaPhiM4D<double>> myP4 {myPt, myEta, 0, myM};
        return myP4.energy();
      }
 
      double rapidity (ObjectId<CI,CM> object) const
      {
        const double myEta = this->eta(object);
        const double myM = this->m(object);
        if (myM == 0)
          return myEta; // for massless particles, rapidity is the same as eta
        const double myPt = this->pt(object);
 
        // not reading phi, not part of the rapidity calculation
        const ROOT::Math::LorentzVector<ROOT::Math::PtEtaPhiM4D<double>> myP4 {myPt, myEta, 0, myM};
        return myP4.Rapidity();
      }
    };
 
 

    template<ContainerIdConcept CIVal,typename CMVal>
    struct CoreMomentumAccessorsPtEtaPhi
    {
      using CI = CIVal;
      using CM = CMVal;
 
      ColumnAccessor<CI,RetypeColumn<double,float>,CM> pt;
      ColumnAccessor<CI,RetypeColumn<double,float>,CM> eta;
      ColumnAccessor<CI,RetypeColumn<double,float>,CM> phi;
 
      CoreMomentumAccessorsPtEtaPhi (ColumnarTool<CM>& columnarTool)
        : pt (columnarTool, "pt"), eta (columnarTool, "eta"), phi (columnarTool, "phi")
      {}

      CoreMomentumAccessorsPtEtaPhi (ColumnarTool<CM>& columnarTool, const std::string& prefix)
        : pt (columnarTool, prefix + "pt"), eta (columnarTool, prefix + "eta"), phi (columnarTool, prefix + "phi")
      {}
    };

    template<ContainerIdConcept CI, typename CM> struct CoreMomentumAccessorsPtEtaPhiReadM
      : CoreMomentumAccessorsPtEtaPhi<CI,CM>
    {
      ColumnAccessor<CI,RetypeColumn<double,float>,CM> m;
 
      CoreMomentumAccessorsPtEtaPhiReadM (ColumnarTool<CM>& columnarTool)
        : CoreMomentumAccessorsPtEtaPhi<CI,CM> (columnarTool), m (columnarTool, "m")
      {}
 
      CoreMomentumAccessorsPtEtaPhiReadM (ColumnarTool<CM>& columnarTool, const std::string& prefix)
        : CoreMomentumAccessorsPtEtaPhi<CI,CM> (columnarTool, prefix), m (columnarTool, prefix + "m")
      {}
    };

    template<ContainerIdConcept CI, typename CM>
    struct CoreMomentumAccessorsPtEtaPhiFixedM : CoreMomentumAccessorsPtEtaPhi<CI,CM>
    {
      double mValue = 0;
 
      double m (ObjectId<CI,CM> /*object*/) const noexcept {
        return mValue; }
 
      CoreMomentumAccessorsPtEtaPhiFixedM (ColumnarTool<CM>& columnarTool, double val_mValue)
        : CoreMomentumAccessorsPtEtaPhi<CI,CM> (columnarTool), mValue (val_mValue)
      {}
    };
 

    template<ContainerIdConcept CI, typename CM = ColumnarModeDefault>
    class IMomentumAccessors
    {
    public:
      virtual ~IMomentumAccessors () = default;
 
      virtual double pt (ObjectId<CI,CM> object) const = 0;
      virtual double eta (ObjectId<CI,CM> object) const = 0;
      virtual double phi (ObjectId<CI,CM> object) const = 0;
      virtual double m (ObjectId<CI,CM> object) const = 0;
      virtual double e (ObjectId<CI,CM> object) const = 0;
      virtual double rapidity (ObjectId<CI,CM> object) const = 0;
    };

    template<typename CoreAccessors>
    class MomentumAccessorsImp : public IMomentumAccessors<typename CoreAccessors::CI,typename CoreAccessors::CM>
    {
    public:
      using CI = CoreAccessors::CI;
      using CM = typename CoreAccessors::CM;
 
      template<typename... Args>
      MomentumAccessorsImp (Args&&... args)
        : m_coreAccessors (std::forward<Args> (args)...) {}
 
      virtual double pt (ObjectId<CI,CM> object) const override {
        return m_coreAccessors.pt (object); }
      virtual double eta (ObjectId<CI,CM> object) const override {
        return m_coreAccessors.eta (object); }
      virtual double phi (ObjectId<CI,CM> object) const override {
        return m_coreAccessors.phi (object); }
      virtual double m (ObjectId<CI,CM> object) const override  {
        return m_coreAccessors.m (object); }
      virtual double e (ObjectId<CI,CM> object) const override {
        return m_coreAccessors.e (object); }
      virtual double rapidity (ObjectId<CI,CM> object) const override {
        return m_coreAccessors.rapidity (object); }
 
    private:
      CoreAccessors m_coreAccessors;
    };
  }

 
  template<ContainerIdConcept CI, typename CM = ColumnarModeDefault>
  class MomentumAccessors final
  {
  public:
 
    MomentumAccessors () noexcept
    {
      reset (std::in_place_type<Detail::MomentumAccessorsIParticle<CI,CM>>);
    }
 
    template<typename MyMomentumAccessors,typename... Args>
    MomentumAccessors (std::in_place_type_t<MyMomentumAccessors>, Args&&... args)
    {
      reset (std::in_place_type<MyMomentumAccessors>, std::forward<Args> (args)...);
    }
 
    template<typename MyMomentumAccessors,typename... Args>
    void reset (std::in_place_type_t<MyMomentumAccessors>, Args&&... args)
    {
      m_accessors = std::make_shared<Detail::MomentumAccessorsImp<MyMomentumAccessors>> (std::forward<Args> (args)...);
    }

    [[nodiscard]] double pt (ObjectId<CI,CM> object) const {
      return m_accessors->pt (object); }
    [[nodiscard]] double eta (ObjectId<CI,CM> object) const {
      return m_accessors->eta (object); }
    [[nodiscard]] double phi (ObjectId<CI,CM> object) const {
      return m_accessors->phi (object); }
    [[nodiscard]] double m (ObjectId<CI,CM> object) const {
      return m_accessors->m (object); }
    [[nodiscard]] double e (ObjectId<CI,CM> object) const {
      return m_accessors->e (object); }
    [[nodiscard]] double rapidity (ObjectId<CI,CM> object) const {
      return m_accessors->rapidity (object); }
 
  private:

    std::shared_ptr<const Detail::IMomentumAccessors<CI,CM>> m_accessors;
  };

  template<ContainerIdConcept CI, typename CM>
  void resetIParticle (MomentumAccessors<CI,CM>& accessors) {
    accessors.reset (std::in_place_type<Detail::MomentumAccessorsIParticle<CI,CM>>); }
  template<ContainerIdConcept CI, typename CM>
  void resetPtEtaPhiReadM (MomentumAccessors<CI,CM>& accessors, ColumnarTool<CM>& columnarTool) {
    accessors.reset (std::in_place_type<Detail::FullMomentumAccessorsPtEtaPhiM<Detail::CoreMomentumAccessorsPtEtaPhiReadM<CI,CM>>>, columnarTool); }
  template<ContainerIdConcept CI, typename CM>
  void resetPtEtaPhiFixedM (MomentumAccessors<CI,CM>& accessors, ColumnarTool<CM>& columnarTool, double mValue) {
    accessors.reset (std::in_place_type<Detail::FullMomentumAccessorsPtEtaPhiM<Detail::CoreMomentumAccessorsPtEtaPhiFixedM<CI,CM>>>, columnarTool, mValue); }
  template<ContainerIdConcept CI, typename CM>
  void resetJet (MomentumAccessors<CI,CM>& accessors, ColumnarTool<CM>& columnarTool) {
    resetPtEtaPhiReadM (accessors, columnarTool); }
  template<ContainerIdConcept CI, typename CM>
  void resetJetConstituentScale (MomentumAccessors<CI,CM>& accessors, ColumnarTool<CM>& columnarTool, const std::string& prefix = "JetConstitScaleMomentum_") {
    accessors.reset (std::in_place_type<Detail::FullMomentumAccessorsPtEtaPhiM<Detail::CoreMomentumAccessorsPtEtaPhiReadM<CI,CM>>>, columnarTool, prefix); }
  template<ContainerIdConcept CI, typename CM>
  void resetElectron (MomentumAccessors<CI,CM>& accessors, ColumnarTool<CM>& columnarTool) {
    resetPtEtaPhiFixedM (accessors, columnarTool, ParticleConstants::electronMassInMeV); }
  template<ContainerIdConcept CI, typename CM>
  void resetPhoton (MomentumAccessors<CI,CM>& accessors, ColumnarTool<CM>& columnarTool) {
    resetPtEtaPhiFixedM (accessors, columnarTool, ParticleConstants::photonMassInMeV); }
  template<ContainerIdConcept CI, typename CM>
  void resetEgamma (MomentumAccessors<CI,CM>& accessors, ColumnarTool<CM>& columnarTool) {
    resetPtEtaPhiReadM (accessors, columnarTool); }
  template<ContainerIdConcept CI, typename CM>
  void resetMuon (MomentumAccessors<CI,CM>& accessors, ColumnarTool<CM>& columnarTool) {
    resetPtEtaPhiFixedM (accessors, columnarTool, ParticleConstants::muonMassInMeV); }
  template<ContainerIdConcept CI, typename CM>
  void resetTau (MomentumAccessors<CI,CM>& accessors, ColumnarTool<CM>& columnarTool) {
    resetPtEtaPhiReadM (accessors, columnarTool); }
 
  template<ContainerIdConcept CI, typename CM>
  void resetObjectType (MomentumAccessors<CI,CM>& accessors, ColumnarTool<CM>& columnarTool, xAODType::ObjectType type)
  {
    switch (type)
    {
      case xAODType::ObjectType::Electron:
        resetElectron (accessors, columnarTool);
        break;
      case xAODType::ObjectType::Photon:
        resetPhoton (accessors, columnarTool);
        break;
      case xAODType::ObjectType::Muon:
        resetMuon (accessors, columnarTool);
        break;
      case xAODType::ObjectType::Tau:
        resetTau (accessors, columnarTool);
        break;
      case xAODType::ObjectType::Jet:
        resetJet (accessors, columnarTool);
        break;
      default:
        throw std::runtime_error ("Unknown object type for momentum accessors: " + std::to_string(static_cast<unsigned>(type)));
    }
  }
}
 
#endif