Hadrons merge

include/DDML/EndcapTriggerTwoAngleBIBAE.h

@@ -28,4 +28,4 @@ class EndcapTriggerTwoAngleBIBAE : public TriggerInterface {
 };
 
 } // namespace ddml
-#endif
+#endif

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/OctogonalBarrelTrigger.h

@@ -28,4 +28,4 @@ class OctogonalBarrelTrigger : public TriggerInterface {
 };
 
 } // namespace ddml
-#endif
+#endif

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() = default;
+
+  virtual ~PionCloudsModel() = default;
+
+  /// 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 {
@@ -26,7 +31,10 @@ class PolyhedraBarrelGeometry : public GeometryInterface {
   /// 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);
   }
 
@@ -50,6 +58,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/TriggerInterface.h

@@ -22,4 +22,4 @@ class TriggerInterface {
 
 } // namespace ddml
 
-#endif
+#endif

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)

scripts/test_onnx.mac

@@ -3,13 +3,24 @@
 #  example script to run inference on a GAN w/ ONNX runtime
 
 # --- the particle gun
-/gps/pos/centre 0 0 0
-/gps/particle  gamma
-#/gps/particle  pi-
-/gps/energy    50 GeV
+
+# At right side face (perpendicular to x=0)
+#/gps/pos/centre 18.47 0 0 #0 180.47 150 #### this is ILD calo Face #120 180.48 120 #30 180.48 0 #30 0 0
+#30 0 0
+
+#/gps/pos/centre 0 0 0 #0 180.47 150 #### this is ILD calo Face #120 180.48 120 #30 180.48 0 #30 0 0
+#30 0 0
+#/gps/particle  gamma
 
 # for SimpleCalo
-/gps/direction  -0.4152 1 0 #(pi/8 + pi/4) #1 0.414213562 0 #(pi/8) #-1 2.414213562 0 #(pi/8 + 1 deg) # 0.25 0.25 0.5  #0 0.5 0.655 # ( 37 deg in theta- most you can get in endcap from IP )
+#/gps/direction  1 0 0 #0 1 0 #-0.4152 1 0 #(pi/8 + pi/4) #1 0.414213562 0 #(pi/8) #-1 2.414213562 0 #(pi/8 + 1 deg) # 0.25 0.25 0.5  #0 0.5 0.655 # ( 37 deg in theta- most you can get in endcap from IP )
+
+/gps/particle  pi+
+/gps/energy    50 GeV
+
+/gps/pos/centre -5 0 -15
+/gps/direction -0.037864591009775746 0.9992828792427412 0 # 90 deg imact for 50 GeV p+
+
 
                 # phi barrel
                 #1 0.414213562 0 #(pi/8)
@@ -64,5 +75,5 @@
 #/gps/direction  0.1 -0.8 .1
 
 
-/run/beamOn 10 #100
+/run/beamOn 2000 #100
 

src/CMakeLists.txt

@@ -7,6 +7,7 @@ set(sources
   Par04ExampleVAE.cc
   RegularGridBIBAEModel.cc
   RegularGridTwoAngleBIBAEModel.cc
+  PionCloudsModel.cc
   OctogonalBarrelTrigger.cc
   EndcapTriggerTwoAngleBIBAE.cc
   CaloCloudsTwoAngleModel.cc
@@ -82,6 +83,7 @@ set(headers
   ${PROJECT_SOURCE_DIR}/include/DDML/Par04ExampleVAE.h
   ${PROJECT_SOURCE_DIR}/include/DDML/RegularGridBIBAEModel.h
   ${PROJECT_SOURCE_DIR}/include/DDML/RegularGridTwoAngleBIBAEModel.h
+  ${PROJECT_SOURCE_DIR}/include/DDML/PionCloudsModel.h
   ${PROJECT_SOURCE_DIR}/include/DDML/PolyhedraBarrelGeometry.h
   ${PROJECT_SOURCE_DIR}/include/DDML/ModelInterface.h
   ${PROJECT_SOURCE_DIR}/include/DDML/InferenceInterface.h

src/LoadHdf5.cc

@@ -19,7 +19,7 @@ void LoadHdf5::initialize() {
   // inputs and TensorDimVecs unused
 
   // Open dataset + dataspace
-  std::string datasetName = "layers";
+  std::string datasetName = "spase_points"; //"events"; //"spase_points"; //"layers";
   H5::DataSet dataset = m_file.openDataSet(datasetName);
   dd4hep::printout(dd4hep::DEBUG, "LoadHdf5::initialize", "Accessed HDF5 dataset");
   H5::DataSpace dataspace = dataset.getSpace();
@@ -28,8 +28,8 @@ void LoadHdf5::initialize() {
   // bool H5::DataSpace::isSimple() const;
 
   // Get dimensions
-  int rank = dataspace.getSimpleExtentNdims();
-  std::vector<hsize_t> dims_out(rank);
+  m_rank = dataspace.getSimpleExtentNdims();
+  std::vector<hsize_t> dims_out(m_rank);
   dataspace.getSimpleExtentDims(dims_out.data(), nullptr);
 
   // Calculate the total number of elements
@@ -39,13 +39,19 @@ void LoadHdf5::initialize() {
 
   m_dimsOut = dims_out;
 
-  assert(rank == 4); // assuming 4 dimensional input
+  assert(m_rank == 3 || m_rank == 4); // ensure either 3 dimensional PionCloud input or 4 dimensional input
+  // if m_rank == 4
   // index 0: shower number
   // index 1, 2, 3: x, y, z cell number
 
-  dd4hep::printout(dd4hep::DEBUG, "LoadHdf5::initialize", "Rank %i", rank);
+  // else if m_rank == 3
+  // index 0: shower number
+  // index 1: number of points
+  // index 2: 4 point dimensions (x, y, z, E) - currently in ILD coordinates
+
+  dd4hep::printout(dd4hep::DEBUG, "LoadHdf5::initialize", "Rank %i", m_rank);
 
-  for (int i = 0, N = rank; i < N; ++i) {
+  for (int i = 0, N = m_rank; i < N; ++i) {
     dd4hep::printout(dd4hep::DEBUG, "LoadHdf5::initialize", "dimension %i: %lu", i, (unsigned long)(m_dimsOut[i]));
   }
 
@@ -71,23 +77,29 @@ void LoadHdf5::runInference(const InputVecs&, const TensorDimVecs&, std::vector<
     m_isInitialized = true;
   }
 
-  // keep track of shower index in file
-  // LoadHdf5::incrementCounter();
-
   // If counter exceeds number of showers in file, reset
   if (m_count > m_totalSize) {
     m_count = 0;
   }
 
   // select shower from library
-  std::vector<float> shower(m_library.begin() + m_count * m_dimsOut[1] * m_dimsOut[2] * m_dimsOut[3],
-                            m_library.begin() + (m_count + 1) * m_dimsOut[1] * m_dimsOut[2] * m_dimsOut[3]);
-
-  // enforce length of shower
-  assert(shower.size() == m_dimsOut[1] * m_dimsOut[2] * m_dimsOut[3]);
-
-  // write output
-  output = std::move(shower);
+  if (m_rank == 4) {
+    std::vector<float> shower(m_library.begin() + m_count * m_dimsOut[1] * m_dimsOut[2] * m_dimsOut[3],
+                              m_library.begin() + (m_count + 1) * m_dimsOut[1] * m_dimsOut[2] * m_dimsOut[3]);
+    assert(shower.size() == m_dimsOut[1] * m_dimsOut[2] * m_dimsOut[3]);
+
+    // write output
+    output = std::move(shower);
+  } else if (m_rank == 3) {
+    std::vector<float> shower(m_library.begin() + m_count * m_dimsOut[1] * m_dimsOut[2],
+                              m_library.begin() + (m_count + 1) * m_dimsOut[1] * m_dimsOut[2]);
+    assert(shower.size() == m_dimsOut[1] * m_dimsOut[2]);
+
+    // write output
+    output = std::move(shower);
+  } else {
+    throw std::runtime_error("Shower library does not have the correct dimensions!");
+  }
 
   ++m_count;
 }