Digicontext: Several improvements

Work needed for restructuring O2DPG MC workflows towards the use of a globally pre-generated collision context. * fixes for picking up vertices from collision contexts * possibility to generate vertices in CollContextTool from CCDB entry * smaller cleanup relates to https://its.cern.ch/jira/browse/O2-3622
AliceO2Group · Nov 11, 2024 · adc8cad · adc8cad
1 parent d0c4891
commit adc8cad
Show file tree

Hide file tree

Showing 5 changed files with 62 additions and 28 deletions.
diff --git a/Common/SimConfig/include/SimConfig/SimConfig.h b/Common/SimConfig/include/SimConfig/SimConfig.h
@@ -83,6 +83,7 @@ struct SimConfigData {
   bool mNoGeant = false;                              // if Geant transport should be turned off (when one is only interested in the generated events)
   bool mIsUpgrade = false;                            // true if the simulation is for Run 5
   std::string mFromCollisionContext = "";             // string denoting a collision context file; If given, this file will be used to determine number of events
+                                                      //
   bool mForwardKine = false;                          // true if tracks and event headers are to be published on a FairMQ channel (for reading by other consumers)
   bool mWriteToDisc = true;                           // whether we write simulation products (kine, hits) to disc
   VertexMode mVertexMode = VertexMode::kDiamondParam; // by default we should use die InteractionDiamond parameter
@@ -177,6 +178,10 @@ class SimConfig
   bool writeToDisc() const { return mConfigData.mWriteToDisc; }
   VertexMode getVertexMode() const { return mConfigData.mVertexMode; }
 
+  // returns the pair of collision context filename as well as event prefix encoded
+  // in the mFromCollisionContext string. Returns empty string if information is not available or set.
+  std::pair<std::string, std::string> getCollContextFilenameAndEventPrefix() const;
+
  private:
   SimConfigData mConfigData; //!
 

diff --git a/Common/SimConfig/src/SimConfig.cxx b/Common/SimConfig/src/SimConfig.cxx
@@ -76,7 +76,7 @@ void SimConfig::initOptions(boost::program_options::options_description& options
     "noGeant", bpo::bool_switch(), "prohibits any Geant transport/physics (by using tight cuts)")(
     "forwardKine", bpo::bool_switch(), "forward kinematics on a FairMQ channel")(
     "noDiscOutput", bpo::bool_switch(), "switch off writing sim results to disc (useful in combination with forwardKine)");
-  options.add_options()("fromCollContext", bpo::value<std::string>()->default_value(""), "Use a pregenerated collision context to infer number of events to simulate, how to embedd them, the vertex position etc. Takes precedence of other options such as \"--nEvents\".");
+  options.add_options()("fromCollContext", bpo::value<std::string>()->default_value(""), "Use a pregenerated collision context to infer number of events to simulate, how to embedd them, the vertex position etc. Takes precedence of other options such as \"--nEvents\". The format is COLLISIONCONTEXTFILE.root[:SIGNALNAME] where SIGNALNAME is the event part in the context which is relevant.");
 }
 
 void SimConfig::determineActiveModules(std::vector<std::string> const& inputargs, std::vector<std::string> const& skippedModules, std::vector<std::string>& activeModules, bool isUpgrade)
@@ -270,6 +270,21 @@ void SimConfig::determineReadoutDetectors(std::vector<std::string> const& active
   }
 }
 
+std::pair<std::string, std::string> SimConfig::getCollContextFilenameAndEventPrefix() const
+{
+  // we decompose the argument to fetch
+  // (a) collision contextfilename
+  // (b) sim prefix to use from the context
+  auto pos = mConfigData.mFromCollisionContext.find(':');
+  std::string collcontextfile{mConfigData.mFromCollisionContext};
+  std::string simprefix{mConfigData.mOutputPrefix};
+  if (pos != std::string::npos) {
+    collcontextfile = mConfigData.mFromCollisionContext.substr(0, pos);
+    simprefix = mConfigData.mFromCollisionContext.substr(pos + 1);
+  }
+  return std::make_pair(collcontextfile, simprefix);
+}
+
 bool SimConfig::resetFromParsedMap(boost::program_options::variables_map const& vm)
 {
   using o2::detectors::DetID;
@@ -333,17 +348,8 @@ bool SimConfig::resetFromParsedMap(boost::program_options::variables_map const&
     mConfigData.mFilterNoHitEvents = true;
   }
   mConfigData.mFromCollisionContext = vm["fromCollContext"].as<std::string>();
-  // we decompose the argument to fetch
-  // (a) collision contextfilename
-  // (b) sim prefix to use from the context
-  auto pos = mConfigData.mFromCollisionContext.find(':');
-  std::string collcontextfile{mConfigData.mFromCollisionContext};
-  std::string simprefix{mConfigData.mOutputPrefix};
-  if (pos != std::string::npos) {
-    collcontextfile = mConfigData.mFromCollisionContext.substr(0, pos);
-    simprefix = mConfigData.mFromCollisionContext.substr(pos + 1);
-  }
-  adjustFromCollContext(collcontextfile, simprefix);
+  auto collcontext_simprefix = getCollContextFilenameAndEventPrefix();
+  adjustFromCollContext(collcontext_simprefix.first, collcontext_simprefix.second);
 
   // analyse vertex options
   if (!parseVertexModeString(vm["vertexMode"].as<std::string>(), mConfigData.mVertexMode)) {

diff --git a/DataFormats/simulation/src/DigitizationContext.cxx b/DataFormats/simulation/src/DigitizationContext.cxx
@@ -578,5 +578,7 @@ DigitizationContext DigitizationContext::extractSingleTimeframe(int timeframeid,
   } catch (std::exception) {
     LOG(warn) << "No such timeframe id in collision context. Returing empty object";
   }
+  // fix number of collisions
+  r.setNCollisions(r.mEventRecords.size());
   return r;
 }
diff --git a/Steer/src/CollisionContextTool.cxx b/Steer/src/CollisionContextTool.cxx
@@ -27,6 +27,7 @@
 #include "CommonUtils/ConfigurableParam.h"
 #include <CCDB/BasicCCDBManager.h>
 #include "DataFormatsParameters/GRPLHCIFData.h"
+#include "SimConfig/SimConfig.h"
 
 //
 // Created by Sandro Wenzel on 13.07.21.
@@ -52,11 +53,12 @@ struct Options {
   bool useexistingkinematics = false;
   bool noEmptyTF = false; // prevent empty timeframes; the first interaction will be shifted backwards to fall within the range given by Options.orbits
   int maxCollsPerTF = -1; // the maximal number of hadronic collisions per TF (can be used to constrain number of collisions per timeframe to some maximal value)
-  bool genVertices = false;         // whether to assign vertices to collisions
   std::string configKeyValues = ""; // string to init config key values
   long timestamp = -1;              // timestamp for CCDB queries
   std::string individualTFextraction = ""; // triggers extraction of individuel timeframe components when non-null
                                            // format is path prefix
+  std::string vertexModeString{"kNoVertex"}; // Vertex Mode; vertices will be assigned to collisions of mode != kNoVertex
+  o2::conf::VertexMode vertexMode = o2::conf::VertexMode::kNoVertex;
 };
 
 enum class InteractionLockMode {
@@ -203,7 +205,9 @@ bool parseOptions(int argc, char* argv[], Options& optvalues)
     "first-orbit", bpo::value<double>(&optvalues.firstFractionalOrbit)->default_value(0), "First (fractional) orbit in the run (HBFUtils.firstOrbit + BC from decimal)")(
     "maxCollsPerTF", bpo::value<int>(&optvalues.maxCollsPerTF)->default_value(-1), "Maximal number of MC collisions to put into one timeframe. By default no constraint.")(
     "noEmptyTF", bpo::bool_switch(&optvalues.noEmptyTF), "Enforce to have at least one collision")(
-    "configKeyValues", bpo::value<std::string>(&optvalues.configKeyValues)->default_value(""), "Semicolon separated key=value strings (e.g.: 'TPC.gasDensity=1;...')")("with-vertices", "Assign vertices to collisions.")("timestamp", bpo::value<long>(&optvalues.timestamp)->default_value(-1L), "Timestamp for CCDB queries / anchoring")(
+    "configKeyValues", bpo::value<std::string>(&optvalues.configKeyValues)->default_value(""), "Semicolon separated key=value strings (e.g.: 'TPC.gasDensity=1;...')")(
+    "with-vertices", bpo::value<std::string>(&optvalues.vertexModeString)->default_value("kNoVertex"), "Assign vertices to collisions. Argument is the vertex mode. Defaults to no vertexing applied")(
+    "timestamp", bpo::value<long>(&optvalues.timestamp)->default_value(-1L), "Timestamp for CCDB queries / anchoring")(
     "extract-per-timeframe", bpo::value<std::string>(&optvalues.individualTFextraction)->default_value(""),
     "Extract individual timeframe contexts. Format required: time_frame_prefix[:comma_separated_list_of_signals_to_offset]");
 
@@ -225,9 +229,8 @@ bool parseOptions(int argc, char* argv[], Options& optvalues)
     if (vm.count("use-existing-kine")) {
       optvalues.useexistingkinematics = true;
     }
-    if (vm.count("with-vertices")) {
-      optvalues.genVertices = true;
-    }
+
+    o2::conf::SimConfig::parseVertexModeString(optvalues.vertexModeString, optvalues.vertexMode);
 
     // fix the first orbit and bunch crossing
     // auto orbitbcpair = parseOrbitAndBC(optvalues.firstIRString);
@@ -277,10 +280,9 @@ int main(int argc, char* argv[])
       LOG(info) << "Fetch bcPattern information from CCDB";
       // fetch the GRP Object
       auto& ccdb = o2::ccdb::BasicCCDBManager::instance();
-      ccdb.setTimestamp(options.timestamp);
       ccdb.setCaching(false);
       ccdb.setLocalObjectValidityChecking(true);
-      auto grpLHC = ccdb.get<o2::parameters::GRPLHCIFData>("GLO/Config/GRPLHCIF");
+      auto grpLHC = ccdb.getForTimeStamp<o2::parameters::GRPLHCIFData>("GLO/Config/GRPLHCIF", options.timestamp);
       LOG(info) << "Fetched injection scheme " << grpLHC->getInjectionScheme() << " from CCDB";
       sampler.setBunchFilling(grpLHC->getBunchFilling());
     } else {
@@ -449,14 +451,32 @@ int main(int argc, char* argv[])
 
   auto numTimeFrames = digicontext.finalizeTimeframeStructure(orbitstart, options.orbitsPerTF);
 
-  if (options.genVertices) {
-    // TODO: offer option taking meanVertex directly from CCDB ! "GLO/Calib/MeanVertex"
-    // sample interaction vertices
+  if (options.vertexMode != o2::conf::VertexMode::kNoVertex) {
+    switch (options.vertexMode) {
+      case o2::conf::VertexMode::kCCDB: {
+        // fetch mean vertex from CCDB
+        auto meanv = o2::ccdb::BasicCCDBManager::instance().getForTimeStamp<o2::dataformats::MeanVertexObject>("GLO/Calib/MeanVertex", options.timestamp);
+        if (meanv) {
+          LOG(info) << "Applying vertexing using CCDB mean vertex " << *meanv;
+          digicontext.sampleInteractionVertices(*meanv);
+        } else {
+          LOG(fatal) << "No vertex available";
+        }
+        break;
+      }
 
-    // init this vertex from CCDB or InteractionDiamond parameter
-    const auto& dparam = o2::eventgen::InteractionDiamondParam::Instance();
-    o2::dataformats::MeanVertexObject meanv(dparam.position[0], dparam.position[1], dparam.position[2], dparam.width[0], dparam.width[1], dparam.width[2], dparam.slopeX, dparam.slopeY);
-    digicontext.sampleInteractionVertices(meanv);
+      case o2::conf::VertexMode::kDiamondParam: {
+        // init this vertex from CCDB or InteractionDiamond parameter
+        const auto& dparam = o2::eventgen::InteractionDiamondParam::Instance();
+        o2::dataformats::MeanVertexObject meanv(dparam.position[0], dparam.position[1], dparam.position[2], dparam.width[0], dparam.width[1], dparam.width[2], dparam.slopeX, dparam.slopeY);
+        LOG(info) << "Applying vertexing using DiamondParam mean vertex " << meanv;
+        digicontext.sampleInteractionVertices(meanv);
+        break;
+      }
+      default: {
+        LOG(error) << "Unknown vertex mode ... Not generating vertices";
+      }
+    }
   }
 
   // we fill QED contributions to the context

diff --git a/run/O2PrimaryServerDevice.h b/run/O2PrimaryServerDevice.h
@@ -138,7 +138,8 @@ class O2PrimaryServerDevice final : public fair::mq::Device
     mPrimGen->SetEvent(&mEventHeader);
 
     // A good moment to couple to collision context
-    auto collContextFileName = mSimConfig.getConfigData().mFromCollisionContext;
+    auto collContextFileName_PrefixPair = mSimConfig.getCollContextFilenameAndEventPrefix();
+    auto collContextFileName = collContextFileName_PrefixPair.first;
     if (collContextFileName.size() > 0) {
       LOG(info) << "Simulation has collission context";
       mCollissionContext = o2::steer::DigitizationContext::loadFromFile(collContextFileName);
@@ -147,7 +148,7 @@ class O2PrimaryServerDevice final : public fair::mq::Device
         LOG(info) << "We found " << vertices.size() << " vertices included ";
 
         // initialize the eventID to collID mapping
-        const auto source = mCollissionContext->findSimPrefix(mSimConfig.getOutPrefix());
+        const auto source = mCollissionContext->findSimPrefix(collContextFileName_PrefixPair.second);
         if (source == -1) {
           LOG(fatal) << "Wrong simulation prefix";
         }