Initialize the inference model early

include/DDML/FastMLShower.h

@@ -100,6 +100,7 @@ class FastMLShower : public dd4hep::sim::Geant4FastSimShowerModel {
   /// "ConstructSDandField()"
   virtual void constructSensitives(dd4hep::sim::Geant4DetectorConstructionContext* ctxt) override {
     this->Geant4FastSimShowerModel::constructSensitives(ctxt);
+    m_fastsimML.initialize();
   }
 
   /// User callback to determine if the model is applicable for the particle
@@ -223,6 +224,8 @@ struct FastMLModel {
 
   FastMLModel() : hitMaker(std::make_unique<Geant4FastHitMakerGlobal>()) {}
 
+  void initialize() { inference.initialize(); }
+
   void declareProperties(dd4hep::sim::Geant4Action* plugin) {
     model.declareProperties(plugin);
     inference.declareProperties(plugin);

include/DDML/GeometryInterface.h

@@ -8,6 +8,10 @@
 
 #include "DDML/DDML.h"
 
+namespace dd4hep::sim {
+class Geant4Action;
+}
+
 class G4FastTrack;
 
 namespace ddml {
@@ -22,18 +26,21 @@ namespace ddml {
  */
 
 template <typename T>
-concept GeometryInterface =
-    requires(const T t, const G4FastTrack& aFastTrack, std::vector<SpacePointVec>& spacepoints) {
-      /** compute local direction in coordinate system that has the z-axis pointing
-       * into the calorimeter, normal to the layers
-       */
-      { t.localDirection(aFastTrack) } -> std::same_as<G4ThreeVector>;
-
-      /** convert the local spacepoints to global spacepoints inside sensitive
-       * volumes
-       */
-      { t.localToGlobal(aFastTrack, spacepoints) } -> std::same_as<void>;
-    };
+concept GeometryInterface = requires(const T t, T mt, const G4FastTrack& aFastTrack,
+                                     std::vector<SpacePointVec>& spacepoints, dd4hep::sim::Geant4Action* plugin) {
+  /** compute local direction in coordinate system that has the z-axis pointing
+   * into the calorimeter, normal to the layers
+   */
+  { t.localDirection(aFastTrack) } -> std::same_as<G4ThreeVector>;
+
+  /** convert the local spacepoints to global spacepoints inside sensitive
+   * volumes
+   */
+  { t.localToGlobal(aFastTrack, spacepoints) } -> std::same_as<void>;
+
+  // /// declareProperties will be called from the FastMLShower constructor
+  { mt.declareProperties(plugin) } -> std::same_as<void>;
+};
 
 } // namespace ddml
 

include/DDML/InferenceInterface.h

@@ -6,6 +6,10 @@
 #include <concepts>
 #include <vector>
 
+namespace dd4hep::sim {
+class Geant4Action;
+}
+
 namespace ddml {
 
 /** The basic concept for running inference with one input vector and one
@@ -16,11 +20,15 @@ namespace ddml {
  */
 
 template <typename T>
-concept InferenceInterface =
-    requires(T t, const InputVecs& inputs, const TensorDimVecs& tensDims, std::vector<float>& output) {
-      /// run the inference model - based on input vector and resized outputvector
-      { t.runInference(inputs, tensDims, output) } -> std::same_as<void>;
-    };
+concept InferenceInterface = requires(T t, const InputVecs& inputs, const TensorDimVecs& tensDims,
+                                      std::vector<float>& output, dd4hep::sim::Geant4Action* plugin) {
+  /// run the inference model - based on input vector and resized outputvector
+  { t.runInference(inputs, tensDims, output) } -> std::same_as<void>;
+  /// initialize also has to be there as this will be called during FastMLShower::constructSensitives
+  { t.initialize() } -> std::same_as<void>;
+  /// declareProperties will be called from the FastMLShower constructor
+  { t.declareProperties(plugin) } -> std::same_as<void>;
+};
 
 } // namespace ddml
 

include/DDML/ModelInterface.h

@@ -6,6 +6,10 @@
 
 #include "DDML/DDML.h"
 
+namespace dd4hep::sim {
+class Geant4Action;
+}
+
 class G4FastTrack;
 
 #include <G4ThreeVector.hh>
@@ -18,22 +22,26 @@ namespace ddml {
  *  @date Mar 2023
  */
 template <typename T>
-concept ModelInterface = requires(T t, G4FastTrack const& aFastTrack, G4ThreeVector const& localDir, InputVecs& inputs,
-                                  TensorDimVecs& tensDims, std::vector<float>& output,
-                                  const std::vector<float>& constOutput, std::vector<SpacePointVec>& spacepoints) {
-  /** 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.) and the direction in the local coordinate system (see
-   * @GeometryInterface)
-   */
-  { t.prepareInput(aFastTrack, localDir, inputs, tensDims, output) } -> std::same_as<void>;
-
-  /** interpreting the model output and create a vector of spacepoints per layer
-   * in local coordinates - with the origin at the entry point into the
-   * calorimeter.
-   */
-  { t.convertOutput(aFastTrack, localDir, constOutput, spacepoints) } -> std::same_as<void>;
-};
+concept ModelInterface =
+    requires(T t, G4FastTrack const& aFastTrack, G4ThreeVector const& localDir, InputVecs& inputs,
+             TensorDimVecs& tensDims, std::vector<float>& output, const std::vector<float>& constOutput,
+             std::vector<SpacePointVec>& spacepoints, dd4hep::sim::Geant4Action* plugin) {
+      /** 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.) and the direction in the local coordinate system (see
+       * @GeometryInterface)
+       */
+      { t.prepareInput(aFastTrack, localDir, inputs, tensDims, output) } -> std::same_as<void>;
+
+      /** interpreting the model output and create a vector of spacepoints per layer
+       * in local coordinates - with the origin at the entry point into the
+       * calorimeter.
+       */
+      { t.convertOutput(aFastTrack, localDir, constOutput, spacepoints) } -> std::same_as<void>;
+
+      /// declareProperties will be called from the FastMLShower constructor
+      { t.declareProperties(plugin) } -> std::same_as<void>;
+    };
 
 } // namespace ddml
 

include/DDML/ONNXInference.h

@@ -48,8 +48,6 @@ class ONNXInference {
   std::vector<const char*> m_inNames;
   std::vector<const char*> m_outNames;
 
-  bool m_isInitialized = false;
-
   /// onxx properties for plugin
   std::string m_modelPath = {};
   int m_profileFlag = 0;

include/DDML/TorchInference.h

@@ -36,8 +36,6 @@ class TorchInference {
   torch::jit::script::Module m_jitModule;
   torch::TensorOptions m_options{};
 
-  bool m_isInitialized = false;
-
   /// torch properties for plugin
   std::string m_modelPath = {};
   int m_profileFlag = 0;

src/ONNXInference.cc

@@ -106,13 +106,7 @@ void ONNXInference::initialize() {
 
 /// run the inference model
 void ONNXInference::runInference(const InputVecs& inputs, const TensorDimVecs& tensDims, std::vector<float>& output) {
-  if (!m_isInitialized) {
-    initialize();
-    m_isInitialized = true;
-  }
-
   // create input tensor object from data values
-
   assert(inputs.size() == tensDims.size());
 
   std::vector<Ort::Value> ort_inputs;

src/TorchInference.cc

@@ -39,11 +39,6 @@ void TorchInference::initialize() {
 void TorchInference::runInference(const InputVecs& inputs, const TensorDimVecs& tensDims, std::vector<float>& output) {
   c10::InferenceMode guard(true);
 
-  if (!m_isInitialized) {
-    initialize();
-    m_isInitialized = true;
-  }
-
   if (dd4hep::printLevel() <= dd4hep::DEBUG) {
     dd4hep::printout(dd4hep::DEBUG, "TorchInference::runInference", " ----- TorchInference::runInference -----");
     dd4hep::printout(dd4hep::DEBUG, "TorchInference::runInference", "# inputs = %f : ", inputs.size());