Hadrons merge

include/DDML/CaloCloudsTwoAngleModel.h

@@ -20,9 +20,9 @@ namespace ddml {
 
 class CaloCloudsTwoAngleModel : public ModelInterface {
 public:
-  CaloCloudsTwoAngleModel(){};
+  CaloCloudsTwoAngleModel() {};
 
-  virtual ~CaloCloudsTwoAngleModel(){};
+  virtual ~CaloCloudsTwoAngleModel() {};
 
   /// declare the proerties needed for the plugin
   void declareProperties(dd4hep::sim::Geant4Action* plugin) {

include/DDML/EndcapGeometry.h

@@ -23,7 +23,7 @@ class EndcapGeometry : public GeometryInterface {
   /// initialize the plugin - after properties have been set
   void initialize();
 
-  virtual ~EndcapGeometry(){};
+  virtual ~EndcapGeometry() {};
 
   /// declare the proerties needed for the plugin
   void declareProperties(dd4hep::sim::Geant4Action* plugin) {

include/DDML/EndcapTriggerTwoAngleBIBAE.h

@@ -18,9 +18,9 @@ namespace ddml {
 
 class EndcapTriggerTwoAngleBIBAE : public TriggerInterface {
 public:
-  EndcapTriggerTwoAngleBIBAE(){};
+  EndcapTriggerTwoAngleBIBAE() {};
 
-  virtual ~EndcapTriggerTwoAngleBIBAE(){};
+  virtual ~EndcapTriggerTwoAngleBIBAE() {};
 
   // check trigger
 

include/DDML/GeometryInterface.h

@@ -22,7 +22,7 @@ namespace ddml {
 
 class GeometryInterface {
 public:
-  virtual ~GeometryInterface(){};
+  virtual ~GeometryInterface() {};
 
   /** compute local direction in coordinate system that has the z-axis pointing
    * into the calorimeter, normal to the layers

include/DDML/InferenceInterface.h

@@ -16,7 +16,7 @@ namespace ddml {
 
 class InferenceInterface {
 public:
-  virtual ~InferenceInterface(){};
+  virtual ~InferenceInterface() {};
 
   /// run the inference model - based on input vector and resized outputvector
   virtual void runInference(const InputVecs& inputs, const TensorDimVecs& tensDims, std::vector<float>& output) = 0;

include/DDML/L2LFlowsModel.h

@@ -18,7 +18,7 @@ namespace ddml {
  */
 class L2LFlowsModel : public ModelInterface {
 public:
-  L2LFlowsModel(){};
+  L2LFlowsModel() {};
 
   virtual ~L2LFlowsModel() = default;
 

include/DDML/LoadHdf5.h

@@ -51,6 +51,9 @@ class LoadHdf5 : public InferenceInterface {
   // shower library dimensions
   std::vector<unsigned long> m_dimsOut{};
 
+  // dimensions of shower library
+  int m_rank;
+
   // properties for plugin
   std::string m_filePath = {};
 

include/DDML/ModelInterface.h

@@ -19,7 +19,7 @@ namespace ddml {
 
 class ModelInterface {
 public:
-  virtual ~ModelInterface(){};
+  virtual ~ModelInterface() {};
 
   /** prepare the input vector and resize the output vector for this model
    *  based on the current FastTrack (e.g. extract kinetic energy and incident

include/DDML/ONNXInference.h

@@ -21,7 +21,7 @@ class ONNXInference : public InferenceInterface {
 public:
   ONNXInference();
 
-  virtual ~ONNXInference(){};
+  virtual ~ONNXInference() {};
 
   /// declare the proerties needed for the plugin
   void declareProperties(dd4hep::sim::Geant4Action* plugin);

include/DDML/OctogonalBarrelTrigger.h

@@ -18,9 +18,9 @@ namespace ddml {
 
 class OctogonalBarrelTrigger : public TriggerInterface {
 public:
-  OctogonalBarrelTrigger(){};
+  OctogonalBarrelTrigger() {};
 
-  virtual ~OctogonalBarrelTrigger(){};
+  virtual ~OctogonalBarrelTrigger() {};
 
   // check trigger
 

include/DDML/Par04ExampleVAE.h

@@ -16,9 +16,9 @@ namespace ddml {
 
 class Par04ExampleVAE : public ModelInterface {
 public:
-  Par04ExampleVAE(){};
+  Par04ExampleVAE() {};
 
-  virtual ~Par04ExampleVAE(){};
+  virtual ~Par04ExampleVAE() {};
 
   /// declare the proerties needed for the plugin
   void declareProperties(dd4hep::sim::Geant4Action* plugin) {

include/DDML/PionCloudsModel.h

@@ -0,0 +1,57 @@
+#ifndef PionClouds_H
+#define PionClouds_H
+
+#include "DDML/FastMLShower.h"
+#include "DDML/ModelInterface.h"
+
+namespace ddml {
+
+/** Class for running a point cloud based ML model for fast shower simulation.
+ *  Assumes a cartesian (x,y) coordinates defining the calorimeter planes (layers) and z the depth
+ *  of the calorimeter.
+ *
+ * Based on BiBAETwoAngleModel.
+ *
+ * Implemented here for the PionClouds model intended for hadron shower simulation (ECAL+HCAL).
+ *
+ *  @author A.Korol, DESY
+ *  @author P. McKeown, CERN
+ *  @date Feb. 2025
+ */
+
+class PionCloudsModel : public ModelInterface {
+public:
+  PionCloudsModel() {};
+
+  virtual ~PionCloudsModel() {};
+
+  /// declare the proerties needed for the plugin
+  void declareProperties(dd4hep::sim::Geant4Action* plugin) {
+    plugin->declareProperty("LatentVectorSize", this->m_latentSize);
+  }
+
+  /** prepare the input vector and resize the output vector for this model
+   *  based on the current FastTrack (e.g. extract kinetic energy and incident
+   *  angles.)
+   */
+  virtual void prepareInput(G4FastTrack const& aFastTrack, G4ThreeVector const& localDir, InputVecs& inputs,
+                            TensorDimVecs& tensDims, std::vector<float>& output);
+
+  /** create a vector of spacepoints per layer interpreting the model output
+   */
+  virtual void convertOutput(G4FastTrack const& aFastTrack, G4ThreeVector const& localDir,
+                             const std::vector<float>& output, std::vector<SpacePointVec>& spacepoints);
+
+private:
+  /// model properties for plugin
+  int m_numPoints = 2600;                 // 4148; //2600; //number of points in the shower
+  size_t m_nDims = 4;                     // Number of dimensions
+  int m_latentSize = 3;                   // number of input features (energy, theta, phi)
+  int m_maxNumElements = m_numPoints * 4; // number of space points in the output multiplied by 4 (x,y,z,energy)
+  // number of layers for ILD is 78: int m_nLayer = 78;
+
+  TensorDimVecs m_tensDims = {{1, 1}, {1, 1}, {1, 1}, {1, 3}};
+};
+
+} // namespace ddml
+#endif

include/DDML/PolyhedraBarrelGeometry.h

@@ -12,6 +12,11 @@ namespace ddml {
  *
  *  @author F.Gaede, DESY
  *  @date Mar 2023
+ *
+ * Addiional option included to support Hadronic shower simulation in ECAL + HCAL
+ *  @author P. McKeown, CERN
+ *  @date Feb 2025
+ *
  */
 
 class PolyhedraBarrelGeometry : public GeometryInterface {
@@ -20,15 +25,18 @@ class PolyhedraBarrelGeometry : public GeometryInterface {
     initialize();
   };
 
-  virtual ~PolyhedraBarrelGeometry(){};
+  virtual ~PolyhedraBarrelGeometry() {};
 
   /// initialize the plugin - after properties have been set
   void initialize();
 
   /// declare the proerties needed for the plugin
   void declareProperties(dd4hep::sim::Geant4Action* plugin) {
     plugin->declareProperty("Detector", this->m_detector);
+    plugin->declareProperty("isHadShower", this->m_isHadShower);
+    plugin->declareProperty("HadDetector", this->m_hadDetector);
     plugin->declareProperty("Symmetry", this->m_nSymmetry);
+    plugin->declareProperty("HadSymmetry", this->m_nHadSymmetry);
     plugin->declareProperty("CorrectForAngles", this->m_correctForAngles);
   }
 
@@ -52,6 +60,9 @@ class PolyhedraBarrelGeometry : public GeometryInterface {
   std::string m_detector = {"EcalBarrel"};
   int m_nSymmetry = 8;
   bool m_correctForAngles = false;
+  bool m_isHadShower; //= true;
+  std::string m_hadDetector = {"HcalBarrel"};
+  int m_nHadSymmetry = m_nSymmetry;
 };
 
 } // namespace ddml

include/DDML/RegularGridBIBAEModel.h

@@ -21,9 +21,9 @@ namespace ddml {
 
 class RegularGridBIBAEModel : public ModelInterface {
 public:
-  RegularGridBIBAEModel(){};
+  RegularGridBIBAEModel() {};
 
-  virtual ~RegularGridBIBAEModel(){};
+  virtual ~RegularGridBIBAEModel() {};
 
   /// declare the proerties needed for the plugin
   void declareProperties(dd4hep::sim::Geant4Action* plugin) {

include/DDML/RegularGridGANModel.h

@@ -16,9 +16,9 @@ namespace ddml {
 
 class RegularGridGANModel : public ModelInterface {
 public:
-  RegularGridGANModel(){};
+  RegularGridGANModel() {};
 
-  virtual ~RegularGridGANModel(){};
+  virtual ~RegularGridGANModel() {};
 
   /// declare the proerties needed for the plugin
   void declareProperties(dd4hep::sim::Geant4Action* plugin) {

include/DDML/RegularGridTwoAngleBIBAEModel.h

@@ -23,9 +23,9 @@ namespace ddml {
 
 class RegularGridTwoAngleBIBAEModel : public ModelInterface {
 public:
-  RegularGridTwoAngleBIBAEModel(){};
+  RegularGridTwoAngleBIBAEModel() {};
 
-  virtual ~RegularGridTwoAngleBIBAEModel(){};
+  virtual ~RegularGridTwoAngleBIBAEModel() {};
 
   /// declare the proerties needed for the plugin
   void declareProperties(dd4hep::sim::Geant4Action* plugin) {

include/DDML/TriggerInterface.h

@@ -15,7 +15,7 @@ namespace ddml {
 
 class TriggerInterface {
 public:
-  virtual ~TriggerInterface(){};
+  virtual ~TriggerInterface() {};
 
   virtual bool check_trigger(const G4FastTrack&) = 0;
 };

scripts/ddsim_steer.py

@@ -394,16 +394,27 @@ def aiDanceTorch(kernel):
     CaloClouds = True
     L2LFlows = False
     old_DD4hep = False  ## use for DD4hep versions/commits before ~ Apr 21st 2023
+    hadrons = False
 
     if ild == True:
         ml_barrel_name = "EcalBarrel"
         ml_barrel_symmetry = 8
         ml_endcap_name = "EcalEndcap"
+
+        ## For hadron shower fast simulation
+        ml_had_barrel_name = "HcalBarrel"
+        ml_had_barrel_symmetry = 8
+        ml_had_endcap_name = "HcalEndcap"
     else:
         ml_barrel_name = "ECalBarrel"
         ml_barrel_symmetry = 12
         ml_endcap_name = "ECalEndcap"
 
+        ## For hadron shower fast simulation is needed
+        ml_had_barrel_name = "HCalBarrel"
+        ml_had_barrel_symmetry = 12
+        ml_had_endcap_name = "HCalEndcap"
+
     if BIBAE == True and Two_Angle == False:
         ml_file = "../models/BIBAE_Full_PP_cut.pt"
         ml_model = "RegularGridBIBAEPolyhedraBarrelTorchModel/BarrelModelTorch"
@@ -459,6 +470,7 @@ def aiDanceTorch(kernel):
     model = DetectorConstruction(kernel, str(ml_model))
 
     ##   # Mandatory model parameters
+    model.isHadShower = False
     model.RegionName = "EcalBarrelRegion"
     model.Detector = ml_barrel_name
     model.Symmetry = ml_barrel_symmetry
@@ -515,16 +527,28 @@ def LoadHdf5(kernel):
     BIBAE = True
     Two_Angle = True
     old_DD4hep = False  ## use for DD4hep versions/commits before ~ Apr 21st 2023
+    hadrons = True
 
     if ild == True:
         ml_barrel_name = "EcalBarrel"
         ml_barrel_symmetry = 8
         ml_endcap_name = "EcalEndcap"
+
+        ## For hadron shower fast simulation
+        ml_had_barrel_name = "HcalBarrel"
+        ml_had_barrel_symmetry = 8
+        ml_had_endcap_name = "HcalEndcap"
+
     else:
         ml_barrel_name = "ECalBarrel"
         ml_barrel_symmetry = 12
         ml_endcap_name = "ECalEndcap"
 
+        ## For hadron shower fast simulation is needed
+        ml_had_barrel_name = "HCalBarrel"
+        ml_had_barrel_symmetry = 12
+        ml_had_endcap_name = "HCalEndcap"
+
     if BIBAE == True and Two_Angle == True:
         ml_file = "../models/photons-E5050A-theta9090A-phi9090-p1.hdf5"
         ml_model = (
@@ -533,6 +557,11 @@ def LoadHdf5(kernel):
         ml_model_1 = "LoadHDF5RegularGridTwoAngleBIBAEModelEndcap/EndcapModelTorch"
         ml_correct_angles = False
 
+    if hadrons == True:
+        ml_model_had = "LoadHDF5PionCloudsPCHadronModelPolyhedraBarrel/BarrelModelTorch"
+        ml_had_file = "../models/PionClouds_50GeV_sp_scaled.h5"
+        ml_correct_angles = False
+
     from g4units import GeV, MeV  # DO NOT REMOVE OR MOVE!!!!! (EXCLAMATION MARK)
     from DDG4 import DetectorConstruction, Geant4, PhysicsList
 
@@ -557,6 +586,8 @@ def LoadHdf5(kernel):
         seq.adopt(sensitives)
 
     # -----------------
+    """
+    ## EM in Barrel
     model = DetectorConstruction(kernel, str(ml_model))
 
     ##   # Mandatory model parameters
@@ -578,7 +609,9 @@ def LoadHdf5(kernel):
 
     model.enableUI()
     seq.adopt(model)
+    """
     # -------------------
+    ## EM in Endcap
     model1 = DetectorConstruction(kernel, str(ml_model_1))
 
     ##   # Mandatory model parameters
@@ -599,12 +632,38 @@ def LoadHdf5(kernel):
 
     model1.enableUI()
     seq.adopt(model1)
+
+    # -------------------
+    ## Hadrons in Barrel
+    modelHad1 = DetectorConstruction(kernel, str(ml_model_had))
+
+    ##   # Mandatory model parameters
+    modelHad1.isHadShower = True
+    modelHad1.RegionName = (
+        "EcalBarrelRegion"  # or "HcalBarrelRegion"  ## hadron model triggers in ecal
+    )
+    modelHad1.Detector = ml_barrel_name
+    modelHad1.HadDetector = ml_had_barrel_name
+    modelHad1.Symmetry = ml_barrel_symmetry
+    modelHad1.HadSymmetry = ml_had_barrel_symmetry
+    modelHad1.Enable = True
+    modelHad1.CorrectForAngles = ml_correct_angles
+    # Energy boundaries are optional: Units are GeV
+    modelHad1.ApplicableParticles = {"pi+"}
+    modelHad1.Etrigger = {"pi+": 10.0 * GeV}  # trigger on lower training threshold
+    modelHad1.FilePath = ml_had_file
+    # model.OptimizeFlag = 1
+    # model.IntraOpNumThreads = 1
+
+    modelHad1.enableUI()
+    seq.adopt(modelHad1)
+
     # -------------------
 
     # Now build the physics list:
     phys = kernel.physicsList()
     ph = PhysicsList(kernel, str("Geant4FastPhysics/FastPhysicsList"))
-    ph.EnabledParticles = ["e+", "e-", "gamma"]
+    ph.EnabledParticles = ["e+", "e-", "gamma", "pi+"]
     ph.BeVerbose = True
     ph.enableUI()
     phys.adopt(ph)