Improvements for CollisionContextTool (WIP)

* important step to create collision contexts
  for all timeframes in one step
* needed to have "pregencollcontext" in O2DPG
  work with embedding

DataFormats/simulation/include/SimulationDataFormat/DigitizationContext.h

@@ -110,6 +110,9 @@ class DigitizationContext
   /// Check collision parts for vertex consistency.
   bool checkVertexCompatibility(bool verbose = false) const;
 
+  /// retrieves collision context for a single timeframe (which may be needed by simulation)
+  DigitizationContext extractSingleTimeframe(int timeframeid, std::vector<int> const& sources_to_offset);
+
   /// function reading the hits from a chain (previously initialized with initSimChains
   /// The hits pointer will be initialized (what to we do about ownership??)
   template <typename T>
@@ -125,8 +128,9 @@ class DigitizationContext
   // apply collision number cuts and potential relabeling of eventID
   void applyMaxCollisionFilter(long startOrbit, long orbitsPerTF, int maxColl);
 
-  // finalize timeframe structure (fixes the indices in mTimeFrameStartIndex)
-  void finalizeTimeframeStructure(long startOrbit, long orbitsPerTF);
+  /// finalize timeframe structure (fixes the indices in mTimeFrameStartIndex)
+  // returns the number of timeframes
+  int finalizeTimeframeStructure(long startOrbit, long orbitsPerTF);
 
   // Sample and fix interaction vertices (according to some distribution). Makes sure that same event ids
   // have to have same vertex, as well as event ids associated to same collision.
@@ -167,7 +171,7 @@ class DigitizationContext
   // for each collision we may record/fix the interaction vertex (to be used in event generation)
   std::vector<math_utils::Point3D<float>> mInteractionVertices;
 
-  // the collision records _with_ QED interleaved;
+  // the collision records **with** QED interleaved;
   std::vector<o2::InteractionTimeRecord> mEventRecordsWithQED;
   std::vector<std::vector<o2::steer::EventPart>> mEventPartsWithQED;
 

DataFormats/simulation/src/DigitizationContext.cxx

@@ -19,6 +19,7 @@
 #include <numeric> // for iota
 #include <MathUtils/Cartesian.h>
 #include <DataFormatsCalibration/MeanVertexObject.h>
+#include <filesystem>
 
 using namespace o2::steer;
 
@@ -196,10 +197,52 @@ o2::parameters::GRPObject const& DigitizationContext::getGRP() const
 
 void DigitizationContext::saveToFile(std::string_view filename) const
 {
+  // checks if the path content of filename exists ... otherwise it is created before creating the ROOT file
+  auto ensure_path_exists = [](std::string_view filename) {
+    try {
+      // Extract the directory path from the filename
+      std::filesystem::path file_path(filename);
+      std::filesystem::path dir_path = file_path.parent_path();
+
+      // Check if the directory path is empty (which means filename was just a name without path)
+      if (dir_path.empty()) {
+        // nothing to do
+        return true;
+      }
+
+      // Create directories if they do not exist
+      if (!std::filesystem::exists(dir_path)) {
+        if (std::filesystem::create_directories(dir_path)) {
+          // std::cout << "Directories created successfully: " << dir_path.string() << std::endl;
+          return true;
+        } else {
+          std::cerr << "Failed to create directories: " << dir_path.string() << std::endl;
+          return false;
+        }
+      }
+      return true;
+    } catch (const std::filesystem::filesystem_error& ex) {
+      std::cerr << "Filesystem error: " << ex.what() << std::endl;
+      return false;
+    } catch (const std::exception& ex) {
+      std::cerr << "General error: " << ex.what() << std::endl;
+      return false;
+    }
+  };
+
+  if (!ensure_path_exists(filename)) {
+    LOG(error) << "Filename contains path component which could not be created";
+    return;
+  }
+
   TFile file(filename.data(), "RECREATE");
-  auto cl = TClass::GetClass(typeid(*this));
-  file.WriteObjectAny(this, cl, "DigitizationContext");
-  file.Close();
+  if (file.IsOpen()) {
+    auto cl = TClass::GetClass(typeid(*this));
+    file.WriteObjectAny(this, cl, "DigitizationContext");
+    file.Close();
+  } else {
+    LOG(error) << "Could not write to file " << filename.data();
+  }
 }
 
 DigitizationContext* DigitizationContext::loadFromFile(std::string_view filename)
@@ -391,13 +434,17 @@ void DigitizationContext::applyMaxCollisionFilter(long startOrbit, long orbitsPe
   mEventParts = newparts;
 }
 
-void DigitizationContext::finalizeTimeframeStructure(long startOrbit, long orbitsPerTF)
+int DigitizationContext::finalizeTimeframeStructure(long startOrbit, long orbitsPerTF)
 {
   mTimeFrameStartIndex = getTimeFrameBoundaries(mEventRecords, startOrbit, orbitsPerTF);
   LOG(info) << "Fixed " << mTimeFrameStartIndex.size() << " timeframes ";
   for (auto p : mTimeFrameStartIndex) {
     LOG(info) << p.first << " " << p.second;
   }
+
+  // we will fix the QED contribution as well
+
+  return mTimeFrameStartIndex.size();
 }
 
 std::unordered_map<int, int> DigitizationContext::getCollisionIndicesForSource(int source) const
@@ -483,3 +530,62 @@ void DigitizationContext::sampleInteractionVertices(o2::dataformats::MeanVertexO
     }
   }
 }
+
+DigitizationContext DigitizationContext::extractSingleTimeframe(int timeframeid, std::vector<int> const& sources_to_offset)
+{
+  // go the right orbit
+  // make a return object
+  DigitizationContext r;
+  r.mSimPrefixes = mSimPrefixes;
+  r.mMuBC = mMuBC;
+  try {
+    auto startend = mTimeFrameStartIndex.at(timeframeid);
+
+    auto startindex = startend.first;
+    auto endindex = startend.second;
+
+    std::copy(mEventRecords.begin() + startindex, mEventRecords.begin() + endindex, std::back_inserter(r.mEventRecords));
+    std::copy(mEventParts.begin() + startindex, mEventParts.begin() + endindex, std::back_inserter(r.mEventParts));
+    std::copy(mInteractionVertices.begin() + startindex, mInteractionVertices.begin() + endindex, std::back_inserter(r.mInteractionVertices));
+
+    // copy QED information between these records
+    // we need to find the index into mEventPartsWithQED just smaller than mEventRecords[startindex];
+
+    // auto startIterQED = std::upper_bound(mEventRecordsWithQED.begin(), mEventRecordsWithQED.end(), mEventRecords[startindex]);
+    // auto endIterQED = std::lower_bound(mEventRecordsWithQED.begin(), mEventRecordsWithQED.end(), mEventRecords[endindex]);
+    // auto startQEDindex = std::distance(startIterQED, mEventRecordsWithQED.begin()) - 1;
+    // auto endQEDindex = std::distance(endIterQED, mEventRecordsWithQED.begin()) - 1;
+
+    // std::copy(mEventRecordsWithQED.begin() + startQEDindex, mEventRecordsWithQED.begin() + endQEDindex, std::back_inserter(r.mEventRecordsWithQED));
+    // std::copy(mEventPartsWithQED.begin() + startQEDindex, mEventPartsWithQED.begin() + endQEDindex, std::back_inserter(r.mEventPartsWithQED));
+
+    // let's assume we want to fix the ids for source = source_id
+    // Then we find the first index that has this source_id and take the corresponding number
+    // as offset. Thereafter we subtract this offset from all known event parts.
+    auto perform_offsetting = [&r](int source_id) {
+      auto indices_for_source = r.getCollisionIndicesForSource(source_id);
+      int minvalue = r.mEventParts.size();
+      for (auto& p : indices_for_source) {
+        if (p.first < minvalue) {
+          minvalue = p.first;
+        }
+      }
+      // now fix them
+      for (auto& p : indices_for_source) {
+        auto index_into_mEventParts = p.second;
+        for (auto& part : r.mEventParts[index_into_mEventParts]) {
+          if (part.sourceID == source_id) {
+            part.entryID -= minvalue;
+          }
+        }
+      }
+    };
+    for (auto source_id : sources_to_offset) {
+      perform_offsetting(source_id);
+    }
+
+  } catch (std::exception) {
+    LOG(warn) << "No such timeframe id in collision context";
+  }
+  return r;
+}

Steer/src/CollisionContextTool.cxx

@@ -53,6 +53,8 @@ struct Options {
   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
 };
 
 enum class InteractionLockMode {
@@ -198,7 +200,10 @@ bool parseOptions(int argc, char* argv[], Options& optvalues)
     "timeframeID", bpo::value<int>(&optvalues.tfid)->default_value(0), "Timeframe id of the first timeframe int this context. Allows to generate contexts for different start orbits")(
     "first-orbit", bpo::value<uint32_t>(&optvalues.firstOrbit)->default_value(0), "First orbit in the run (HBFUtils.firstOrbit)")(
     "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");
+    "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")(
+    "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]");
 
   options.add_options()("help,h", "Produce help message.");
 
@@ -431,7 +436,7 @@ int main(int argc, char* argv[])
   // apply max collision per timeframe filters + reindexing of event id (linearisation and compactification)
   digicontext.applyMaxCollisionFilter(options.tfid * options.orbitsPerTF, options.orbitsPerTF, options.maxCollsPerTF);
 
-  digicontext.finalizeTimeframeStructure(options.tfid * options.orbitsPerTF, options.orbitsPerTF);
+  auto numTimeFrames = digicontext.finalizeTimeframeStructure(options.tfid * options.orbitsPerTF, options.orbitsPerTF);
 
   if (options.genVertices) {
     // TODO: offer option taking meanVertex directly from CCDB ! "GLO/Calib/MeanVertex"
@@ -456,5 +461,72 @@ int main(int argc, char* argv[])
   }
   digicontext.saveToFile(options.outfilename);
 
+  // extract individual timeframes
+  if (options.individualTFextraction.size() > 0) {
+    // we are asked to extract individual timeframe components
+
+    LOG(info) << "Extracting individual timeframe collision contexts";
+    // extract prefix path to store these collision contexts
+    // Function to check the pattern and extract tokens from b
+    auto check_and_extract_tokens = [](const std::string& input, std::vector<std::string>& tokens) {
+      // the regular expression pattern for expected input format
+      const std::regex pattern(R"(^([a-zA-Z0-9]+)(:([a-zA-Z0-9]+(,[a-zA-Z0-9]+)*))?$)");
+      std::smatch matches;
+
+      // Check if the input matches the pattern
+      if (std::regex_match(input, matches, pattern)) {
+        // Clear any existing tokens in the vector
+        tokens.clear();
+
+        // matches[1] contains the part before the colon which we save first
+        tokens.push_back(matches[1].str());
+        // matches[2] contains the comma-separated list
+        std::string b = matches[2].str();
+        std::regex token_pattern(R"([a-zA-Z0-9]+)");
+        auto tokens_begin = std::sregex_iterator(b.begin(), b.end(), token_pattern);
+        auto tokens_end = std::sregex_iterator();
+
+        // Iterate over the tokens and add them to the vector
+        for (std::sregex_iterator i = tokens_begin; i != tokens_end; ++i) {
+          tokens.push_back((*i).str());
+        }
+        return true;
+      }
+      LOG(error) << "Argument for --extract-per-timeframe does not match specification";
+      return false;
+    };
+
+    std::vector<std::string> tokens;
+    if (check_and_extract_tokens(options.individualTFextraction, tokens)) {
+      auto path_prefix = tokens[0];
+      std::vector<int> sources_to_offset{};
+
+      LOG(info) << "PREFIX is " << path_prefix;
+
+      for (int i = 1; i < tokens.size(); ++i) {
+        LOG(info) << "Offsetting " << tokens[i];
+        sources_to_offset.push_back(digicontext.findSimPrefix(tokens[i]));
+      }
+
+      // now we are ready to loop over all timeframes
+      for (int tf_id = 0; tf_id < numTimeFrames; ++tf_id) {
+        auto copy = digicontext.extractSingleTimeframe(tf_id, sources_to_offset);
+
+        // each individual case gets QED interactions injected
+        // This should probably be done inside the extraction itself
+        if (digicontext.isQEDProvided()) {
+          auto qedSpec = parseInteractionSpec(options.qedInteraction, ispecs, options.useexistingkinematics);
+          copy.fillQED(qedSpec.name, qedSpec.mcnumberasked, qedSpec.interactionRate);
+        }
+
+        std::stringstream str;
+        str << path_prefix << (tf_id + 1) << "/collisioncontext.root";
+        copy.saveToFile(str.str());
+        LOG(info) << "----";
+        copy.printCollisionSummary(options.qedInteraction.size() > 0);
+      }
+    }
+  }
+
   return 0;
 }