ElectronPhotonVariableNFCorrectionTool.h

Go to the documentation of this file.

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

 * @author Katerina Kazakova <katerina.kazakova@cern.ch>
 * @date   February 2026
 */
 
 
//ATLAS includes
#include "AsgTools/AsgTool.h"
#include "AsgTools/ToolHandleArray.h"
#include <AsgTools/PropertyWrapper.h>
#include "EgammaAnalysisInterfaces/IElectronPhotonShowerShapeFudgeTool.h"
 
#include "AthOnnxInterfaces/IOnnxRuntimeInferenceTool.h"
#include "AthContainers/AuxElement.h"
 
#include "AsgDataHandles/ReadHandleKey.h"
#include "AsgDataHandles/ReadHandle.h"
#include "xAODEventInfo/EventInfo.h"
 
//EDM includes
#include "xAODEgamma/Electron.h"
#include "xAODEgamma/Photon.h"
 
#include <memory>
#include <vector>
#include <string>
 
// ===========================================================================
// Class ElectronPhotonVariableNFCorrectionTool
// ===========================================================================
 
class ElectronPhotonVariableNFCorrectionTool : public asg::AsgTool, virtual public IElectronPhotonShowerShapeFudgeTool
{
 
ASG_TOOL_CLASS(ElectronPhotonVariableNFCorrectionTool, IElectronPhotonShowerShapeFudgeTool)
 
public:
    ElectronPhotonVariableNFCorrectionTool(const std::string& name);

    ~ElectronPhotonVariableNFCorrectionTool() {};

    virtual StatusCode initialize() override;

    virtual const CP::CorrectionCode applyCorrection(xAOD::Photon& photon) const override;

    virtual const CP::CorrectionCode applyCorrection(xAOD::Electron& electron) const override;

    virtual const CP::CorrectionCode correctedCopy(const xAOD::Photon& in_photon,
                                                    xAOD::Photon*& out_photon) const override;

    virtual const CP::CorrectionCode correctedCopy(const xAOD::Electron& in_electron,
                                                    xAOD::Electron*& out_electron) const override;
 
 
private:
    Gaudi::Property<std::string> m_configFile {this, "ConfigFile", "", "The configuration file for Normalizing Flows to use"};
    Gaudi::Property<int> m_nFoldsOverride {this, "NFoldsOverride", -1, "If > 0, override the number of folds to use instead of the value from the config file (must not exceed NFolds in the config)"};
    Gaudi::Property<std::string> m_applyToStr {this, "ApplyTo", "TruthPhotons", "TruthPhotons or All"};
    Gaudi::Property<float>       m_pTcutMeV   {this, "pTcut", 10000.f, "Min photon pT in MeV"};
    
    int m_nFolds{0};

    std::string m_onnxPattern;
    
    SG::ReadHandleKey<xAOD::EventInfo> m_eventInfoKey{this, "EventInfoKey", "EventInfo", "EventInfo key"};
    
    enum class ApplyToMode { TruthPhotons, All };
    ApplyToMode m_applyToMode{ApplyToMode::TruthPhotons};

    bool m_applyShowerShapeCuts{true};

    bool passSelectionCuts(const xAOD::Photon& photon, const std::vector<float>& ss) const;
    
    enum class FoldStrategy { Unknown, EventNumber, EventNumberPhi };
    
    FoldStrategy m_foldStrategy{FoldStrategy::EventNumber};
    
    int selectFold(unsigned long long eventNumber, float phi) const;

    FoldStrategy parseFoldStrategy(const std::string& s) const;
    
    ToolHandleArray<AthOnnx::IOnnxRuntimeInferenceTool> m_onnxToolsForward{this, "OnnxInferenceToolsForward", {}, "Forward ONNX tools per fold"};
    ToolHandleArray<AthOnnx::IOnnxRuntimeInferenceTool> m_onnxToolsBackward{this, "OnnxInferenceToolsBackward", {}, "Backward ONNX tools per fold"};

    static const std::vector<std::string> s_ssVarNames;

    static const std::vector<xAOD::EgammaParameters::ShowerShapeType> s_ssEnums;
    
    struct SSAccessors {
        std::unique_ptr<SG::AuxElement::Accessor<float>> original;
    };

    std::vector<SSAccessors> m_accessors;
 
}; // end class ElectronPhotonVariableNFCorrectionTool
 
#endif