Merge branch 'FastCaloSim_EnergyAndHitGAN' into '21.0'

FastCaloSimGANsimulation

See merge request atlas/athena!31744

Simulation/ISF/ISF_FastCaloSim/ISF_FastCaloSimEvent/CMakeLists.txt

@@ -23,6 +23,8 @@ atlas_depends_on_subdirs(
 find_package( CLHEP )
 find_package( HepPDT )
 find_package( ROOT COMPONENTS Core Tree MathCore Hist RIO Matrix Physics )
+find_package( lwtnn REQUIRED )
+find_package( LibXml2 )
 
 # Remove the --as-needed linker flags:
 atlas_disable_as_needed()
@@ -64,6 +66,7 @@ atlas_add_root_dictionary( ISF_FastCaloSimEvent _dictSource
                            ISF_FastCaloSimEvent/TFCSEnergyParametrization.h 
                            ISF_FastCaloSimEvent/TFCSPCAEnergyParametrization.h 
                            ISF_FastCaloSimEvent/TFCSEnergyBinParametrization.h
+                           ISF_FastCaloSimEvent/TFCSEnergyAndHitGAN.h
                            ISF_FastCaloSimEvent/TFCSLateralShapeParametrization.h
                            ISF_FastCaloSimEvent/TFCSLateralShapeParametrizationHitBase.h
                            ISF_FastCaloSimEvent/TFCSLateralShapeParametrizationHitChain.h
@@ -92,9 +95,9 @@ atlas_add_root_dictionary( ISF_FastCaloSimEvent _dictSource
 atlas_add_library( ISF_FastCaloSimEvent
    ISF_FastCaloSimEvent/*.h ${_dictSource} src/*.cxx
    PUBLIC_HEADERS ISF_FastCaloSimEvent
-   INCLUDE_DIRS ${CLHEP_INCLUDE_DIRS} ${HEPPDT_INCLUDE_DIRS} ${ROOT_INCLUDE_DIRS}
+   INCLUDE_DIRS ${CLHEP_INCLUDE_DIRS} ${HEPPDT_INCLUDE_DIRS} ${ROOT_INCLUDE_DIRS} ${LIBXML2_INCLUDE_DIR}
    DEFINITIONS ${CLHEP_DEFINITIONS}
-   LINK_LIBRARIES ${CLHEP_LIBRARIES} ${HEPPDT_LIBRARIES} ${ROOT_LIBRARIES} AthContainers SGTools AthenaKernel AthenaBaseComps CaloDetDescrLib
+   LINK_LIBRARIES ${CLHEP_LIBRARIES} ${HEPPDT_LIBRARIES} ${ROOT_LIBRARIES} AthContainers SGTools AthenaKernel AthenaBaseComps CaloDetDescrLib ${LWTNN_LIBRARIES} ${LIBXML2_LIBRARIES}
    CaloGeoHelpers TileSimEvent
    PRIVATE_LINK_LIBRARIES GaudiKernel )
 

Simulation/ISF/ISF_FastCaloSim/ISF_FastCaloSimEvent/ISF_FastCaloSimEvent/LinkDef.h

 /*
-  Copyright (C) 2002-2019 CERN for the benefit of the ATLAS collaboration
+  Copyright (C) 2002-2020 CERN for the benefit of the ATLAS collaboration
 */
 
 #include "ISF_FastCaloSimEvent/DoubleArray.h"
@@ -39,6 +39,7 @@
 #include "ISF_FastCaloSimEvent/TFCSEnergyParametrization.h"
 #include "ISF_FastCaloSimEvent/TFCSPCAEnergyParametrization.h"
 #include "ISF_FastCaloSimEvent/TFCSEnergyBinParametrization.h"
+#include "ISF_FastCaloSimEvent/TFCSEnergyAndHitGAN.h"
 
 #include "ISF_FastCaloSimEvent/TFCSLateralShapeParametrization.h"
 #include "ISF_FastCaloSimEvent/TFCSLateralShapeParametrizationHitBase.h"
@@ -202,6 +203,7 @@
 #pragma link C++ class TFCSEnergyParametrization+;
 #pragma link C++ class TFCSPCAEnergyParametrization-;
 #pragma link C++ class TFCSEnergyBinParametrization+;
+#pragma link C++ class TFCSEnergyAndHitGAN-;
 
 #pragma link C++ class TFCSLateralShapeParametrization+;
 #pragma link C++ class TFCSLateralShapeParametrizationHitBase+;

Simulation/ISF/ISF_FastCaloSim/ISF_FastCaloSimEvent/ISF_FastCaloSimEvent/TFCSEnergyAndHitGAN.h

+/*
+  Copyright (C) 2002-2020 CERN for the benefit of the ATLAS collaboration
+*/
+
+#ifndef ISF_FASTCALOSIMEVENT_TFCSEnergyAndHitGAN_h
+#define ISF_FASTCALOSIMEVENT_TFCSEnergyAndHitGAN_h
+
+#include "ISF_FastCaloSimEvent/TFCSParametrizationBinnedChain.h"
+#include "ISF_FastCaloSimEvent/TFCSSimulationState.h"
+#include "lwtnn/LightweightGraph.hh"
+#include <string>
+#include "TH2D.h"
+
+class TFCSEnergyAndHitGAN:public TFCSParametrizationBinnedChain {
+public:
+  // type of input requested by lwtnn
+  typedef std::map<std::string, std::map<std::string, double> >  NetworkInputs ;
+  typedef std::map<std::string, double> NetworkOutputs;
+  typedef std::map<int, TH2D> Binning;
+
+  TFCSEnergyAndHitGAN(const char* name=nullptr, const char* title=nullptr);
+  virtual ~TFCSEnergyAndHitGAN();
+
+  ///Status bit for chain persistency
+  enum FCSGANfreemem {
+     kGANfreemem = BIT(17) ///< Set this bit in the TObject bit if the memory for m_input should be freed after reading in athena
+  };
+
+  bool GANfreemem() const {return TestBit(kGANfreemem);};
+  void set_GANfreemem() {SetBit(kGANfreemem);};
+  void reset_GANfreemem() {ResetBit(kGANfreemem);};
+
+  /// use the layer to be done as binning of the GAN chain
+  virtual int get_bin(TFCSSimulationState& simulstate,const TFCSTruthState*, const TFCSExtrapolationState*) const override {return simulstate.getAuxInfo<int>("GANlayer"_FCShash);};
+  virtual const std::string get_variable_text(TFCSSimulationState& simulstate,const TFCSTruthState*, const TFCSExtrapolationState*) const override;
+  
+  unsigned int get_nr_of_init(unsigned int bin) const;
+  void set_nr_of_init(unsigned int bin,unsigned int ninit);
+  
+  const Binning& get_Binning() const {return m_Binning;};
+  const lwt::LightweightGraph* get_graph() const {return m_graph;};
+  const std::string* get_input() const {return m_input;};
+  
+  bool initializeNetwork(int pid,int etaMin,std::string FastCaloGANInputFolderName);
+  
+  bool fillFastCaloGanNetworkInputs(TFCSSimulationState& simulstate,const TFCSTruthState* truth, NetworkInputs & inputs,double & trueEnergy) const;
+  bool fillEnergy(TFCSSimulationState& simulstate, const TFCSTruthState* truth, const TFCSExtrapolationState* extrapol, NetworkInputs inputs) const;
+  virtual FCSReturnCode simulate(TFCSSimulationState& simulstate,const TFCSTruthState* truth, const TFCSExtrapolationState* extrapol) const override;
+  
+  virtual void Print(Option_t *option="") const override;
+
+  static void unit_test(TFCSSimulationState* simulstate=nullptr,const TFCSTruthState* truth=nullptr, const TFCSExtrapolationState* extrapol=nullptr);
+
+protected:  
+  void GetBinning(int pid,int etaMax,std::string FastCaloGANInputFolderName);
+  
+private:
+  std::vector< int > m_bin_ninit;
+  
+  //Persistify configuration in string m_input. A custom Streamer(...) builds m_graph on the fly when reading from file.
+  //Inside Athena, if GANfreemem() is true, the content of m_input is deleted after reading in order to free memory
+  std::string*     m_input=nullptr;
+  lwt::LightweightGraph* m_graph=nullptr;//!Do not persistify
+  
+  Binning m_Binning;
+  
+  // specific to architecture
+  // preprocessing of input
+  int m_GANLatentSize = 0;
+
+  ClassDefOverride(TFCSEnergyAndHitGAN,2)  //TFCSEnergyAndHitGAN
+};
+
+#endif