diff --git a/DataFormats/simulation/include/SimulationDataFormat/DigitizationContext.h b/DataFormats/simulation/include/SimulationDataFormat/DigitizationContext.h index f531cf6e7f870..98b393a058522 100644 --- a/DataFormats/simulation/include/SimulationDataFormat/DigitizationContext.h +++ b/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 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 @@ -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> mInteractionVertices; - // the collision records _with_ QED interleaved; + // the collision records **with** QED interleaved; std::vector mEventRecordsWithQED; std::vector> mEventPartsWithQED; diff --git a/DataFormats/simulation/src/DigitizationContext.cxx b/DataFormats/simulation/src/DigitizationContext.cxx index f3f40d77042a5..ff71cc4eafbc3 100644 --- a/DataFormats/simulation/src/DigitizationContext.cxx +++ b/DataFormats/simulation/src/DigitizationContext.cxx @@ -19,6 +19,7 @@ #include // for iota #include #include +#include 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 DigitizationContext::getCollisionIndicesForSource(int source) const @@ -483,3 +530,62 @@ void DigitizationContext::sampleInteractionVertices(o2::dataformats::MeanVertexO } } } + +DigitizationContext DigitizationContext::extractSingleTimeframe(int timeframeid, std::vector 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; +} diff --git a/Steer/src/CollisionContextTool.cxx b/Steer/src/CollisionContextTool.cxx index db61343ddf586..5d7fd7d7d345c 100644 --- a/Steer/src/CollisionContextTool.cxx +++ b/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(&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(&optvalues.firstOrbit)->default_value(0), "First orbit in the run (HBFUtils.firstOrbit)")( "maxCollsPerTF", bpo::value(&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(&optvalues.configKeyValues)->default_value(""), "Semicolon separated key=value strings (e.g.: 'TPC.gasDensity=1;...')")("with-vertices", "Assign vertices to collisions.")("timestamp", bpo::value(&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(&optvalues.configKeyValues)->default_value(""), "Semicolon separated key=value strings (e.g.: 'TPC.gasDensity=1;...')")("with-vertices", "Assign vertices to collisions.")("timestamp", bpo::value(&optvalues.timestamp)->default_value(-1L), "Timestamp for CCDB queries / anchoring")( + "extract-per-timeframe", bpo::value(&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& 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 tokens; + if (check_and_extract_tokens(options.individualTFextraction, tokens)) { + auto path_prefix = tokens[0]; + std::vector 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; }