egammaTransformerCalibTool.h

Go to the documentation of this file.

/*
  Copyright (C) 2002-2026 CERN for the benefit of the ATLAS collaboration
*/
 
#ifndef EGAMMATRANSFORMERCALIB_EGAMMATRANSFORMERCALIBTOOL_H
#define EGAMMATRANSFORMERCALIB_EGAMMATRANSFORMERCALIBTOOL_H
 
// Package includes
#include "EgammaAnalysisInterfaces/IegammaMVACalibTool.h"
#include "FlavorTagInference/SaltModel.h"
#include "xAODEgamma/EgammaEnums.h"
#include "egammaLayerRecalibTool/egammaLayerRecalibTool.h"
#include "egammaInterfaces/IegammaCellRecoveryTool.h"
#include "xAODEgamma/EgammaxAODHelpers.h"
 
// Framework includes
#include "AsgTools/AsgTool.h"
#include "AsgTools/PropertyWrapper.h"
 
// Root includes
#include "TFormula.h"
#include "TH2Poly.h"
#include "TObject.h"
#include "TString.h"
 
// STL includes
#include <functional>
#include <memory>
#include <string>

 
class egammaTransformerCalibTool : public asg::AsgTool,
                                   virtual public IegammaMVACalibTool {
  ASG_TOOL_CLASS(egammaTransformerCalibTool, IegammaMVACalibTool)
 public:
  egammaTransformerCalibTool(const std::string& type);
  virtual ~egammaTransformerCalibTool() override;
 
  virtual StatusCode initialize() override;

  float getEnergy(const xAOD::CaloCluster& clus, const xAOD::Egamma* eg,
                  const egammaMVACalib::GlobalEventInfo& gei =
                      egammaMVACalib::GlobalEventInfo()) const override final;
 
 private:
  Gaudi::Property<int> m_particleType{
      this, "ParticleType", xAOD::EgammaParameters::NumberOfEgammaTypes,
      "What type of particle do we use"};
 
  Gaudi::Property<bool> m_clusterEif0{
      this, "useClusterIf0", true, "Use cluster energy if MVA response is 0"};

  Gaudi::Property<std::string> m_folder{this, "folder", "",
                                        "string with folder for weight files"};

  Gaudi::Property<bool> m_isMC{this, "isMC", false, "Whether the input file is MC or data"};
 
  Gaudi::Property<bool> m_useLayerCorrected{this, "useLayerCorrection", true,
                                            "whether to use layer corrections"};
 
  Gaudi::Property<std::string> m_layerCalibTune{this, "layerRecalibrationTune", "es2025_run3_extrapolate_gnn_v0",
                                            "layer to use for layer corrections"};
 
  Gaudi::Property<bool> m_useSaccCorrection{this, "useSaccCorrection", true,
                                            "whether to use SACC correction for layer recalibration"};
 
  Gaudi::Property<bool> m_useFixForMissingCells{this, "useFixForMissingCells", true,
                                            "whether to apply fix for missing cells in layer recalibration, this is applied by default, this must be set to true whenever the timing cut fix is applied. Otherwise the layer calibration will not have correct scale factors for cell energies after timing cut fix"};
 
  Gaudi::Property<bool> m_useExtraLayerScales{this, "useExtraLayerScales", false,
                                            "whether to apply extra layer scales, this is for systematics studies, by default it is false and the extra layer scales are set to 1.0"};
 
  Gaudi::Property<std::string> m_electronModelFile{
      this, "ElectronModelFile", "electron_model_calibration.onnx",
      "ONNX file for electron transformer model"};
 
  Gaudi::Property<std::string> m_unconvertedPhotonModelFile{
      this, "UnconvertedPhotonModelFile",
        "unconverted_photon_model_calibration.onnx",
      "ONNX file for unconverted photon transformer model"};
 
  Gaudi::Property<std::string> m_convertedPhotonModelFile{
      this, "ConvertedPhotonModelFile",
      "converted_photon_model_calibration.onnx",
      "ONNX file for converted photon transformer model"};
 
  Gaudi::Property<std::string> m_forwardElectronModelFile{
      this, "ForwardElectronModelFile",
      "",
      "ONNX file for forward electron transformer model, this is NOT "
      "implemented yet"};
 
  // Transformer model utils
  std::unique_ptr<const FlavorTagInference::SaltModel> m_saltModel;
  int m_num_cluster_features = 0;
  int m_num_cell_features = 0;
  static constexpr double m_timeCut = 12.;
 
  std::unique_ptr<egammaLayerRecalibTool> m_layerRecalibTool = nullptr;
 
  // LG: scale for MG and LG need to be added later since the amount of work is not small.

  ToolHandle<IegammaCellRecoveryTool> m_egammaCellRecoveryTool{
    this,
    "egammaCellRecoveryTool",
    "",
    "Optional tool that adds cells in L2 or L3 "
    "that could have been rejected by timing cut"
  };

  StatusCode setupTransformerModel(const std::string& fileName);
 
};
 
#endif