TPCLoopers.cxx

Go to the documentation of this file.

// Copyright 2024-2025 CERN and copyright holders of ALICE O2.
// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
// All rights not expressly granted are reserved.
//
// This software is distributed under the terms of the GNU General Public
// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
//
// In applying this license CERN does not waive the privileges and immunities
// granted to it by virtue of its status as an Intergovernmental Organization
// or submit itself to any jurisdiction.
 
 
#include "Generators/TPCLoopers.h"
#include "CCDB/CCDBTimeStampUtils.h"
#include "CCDB/CcdbApi.h"
#include "DetectorsRaw/HBFUtils.h"
#include <SimConfig/SimConfig.h>
#include "TF1.h"
#include <filesystem>
#include <algorithm>
#include <limits>
#include <memory>
#include <SimulationDataFormat/ParticleStatus.h>
#include "SimulationDataFormat/MCGenProperties.h"
#include <iostream>
#include <fstream>
#include "TDatabasePDG.h"
 
// Static Ort::Env instance for multiple onnx model loading
Ort::Env global_env(ORT_LOGGING_LEVEL_WARNING, "GlobalEnv");
 
// This class is responsible for loading the scaler parameters from a JSON file
// and applying the inverse transformation to the generated data.
 
void Scaler::load(const std::string& filename)
{
  std::ifstream file(filename);
  if (!file.is_open()) {
    throw std::runtime_error("Error: Could not open scaler file!");
  }
 
  std::string json_str((std::istreambuf_iterator<char>(file)), std::istreambuf_iterator<char>());
  file.close();
 
  rapidjson::Document doc;
  doc.Parse(json_str.c_str());
 
  if (doc.HasParseError()) {
    throw std::runtime_error("Error: JSON parsing failed!");
  }
 
  normal_min = jsonArrayToVector(doc["normal"]["min"]);
  normal_max = jsonArrayToVector(doc["normal"]["max"]);
  outlier_center = jsonArrayToVector(doc["outlier"]["center"]);
  outlier_scale = jsonArrayToVector(doc["outlier"]["scale"]);
}
 
std::vector<double> Scaler::inverse_transform(const std::vector<double>& input)
{
  std::vector<double> output;
  for (int i = 0; i < input.size(); ++i) {
    if (i < input.size() - 2) {
      output.push_back(input[i] * (normal_max[i] - normal_min[i]) + normal_min[i]);
    } else {
      output.push_back(input[i] * outlier_scale[i - (input.size() - 2)] + outlier_center[i - (input.size() - 2)]);
    }
  }
 
  return output;
}
 
std::vector<double> Scaler::jsonArrayToVector(const rapidjson::Value& jsonArray)
{
  std::vector<double> vec;
  for (int i = 0; i < jsonArray.Size(); ++i) {
    vec.push_back(jsonArray[i].GetDouble());
  }
  return vec;
}
 
// This class loads the ONNX model and generates samples using it.
 
ONNXGenerator::ONNXGenerator(Ort::Env& shared_env, const std::string& model_path)
  : env(shared_env), session(nullptr)
{
  Ort::SessionOptions session_options;
  session_options.SetIntraOpNumThreads(1);
  session = Ort::Session(env, model_path.c_str(), session_options);
}
 
std::vector<double> ONNXGenerator::generate_sample()
{
  Ort::AllocatorWithDefaultOptions allocator;
 
  // Generate a latent vector (z)
  std::vector<float> z(100);
  for (auto& v : z) {
    v = rand_gen.Gaus(0.0, 1.0);
  }
 
  // Prepare input tensor
  std::vector<int64_t> input_shape = {1, 100};
  // Get memory information
  Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
 
  // Create input tensor correctly
  Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
    memory_info, z.data(), z.size(), input_shape.data(), input_shape.size());
  // Run inference
  const char* input_names[] = {"z"};
  const char* output_names[] = {"output"};
  auto output_tensors = session.Run(Ort::RunOptions{nullptr}, input_names, &input_tensor, 1, output_names, 1);
 
  // Extract output
  float* output_data = output_tensors.front().GetTensorMutableData<float>();
  // Get the size of the output tensor
  auto output_tensor_info = output_tensors.front().GetTensorTypeAndShapeInfo();
  size_t output_data_size = output_tensor_info.GetElementCount(); // Total number of elements in the tensor
  std::vector<double> output;
  for (int i = 0; i < output_data_size; ++i) {
    output.push_back(output_data[i]);
  }
 
  return output;
}
 
namespace o2
{
namespace eventgen
{
 
namespace
{
// Radial limits of the region from which a looper can still reach the TPC
// sensitive gas. The field cage positions are those used by the
// "ExcludeFCGap" selection in o2::tpc::Detector::ProcessHits()
// A looper can enter the sensitive gas from just outside it, so an additional margin is set
// with a factor two over the largest radial excursion which was observed in validation (~4.2 cm).
//
// No cut is applied on z because loopers spiral the field lines and a vertex as far as |z| = 283 cm
// feeds hits into the gas. A cut here would discard loopers that produce TPC signals.
constexpr double kFcLxIn = 82.428409;      // cm, inner field cage strips
constexpr double kRodROut = 254.25 + 2.2;  // cm, outer field cage rods plus their radial size
constexpr double kLooperRadialReach = 10.; // cm, margin for the helix sweep
constexpr double kTPCActiveRMin = kFcLxIn - kLooperRadialReach;
constexpr double kTPCActiveRMax = kRodROut + kLooperRadialReach;
} // namespace
 
bool GenTPCLoopers::isInTPCActiveVolume(double vx, double vy) const
{
  const double vt = std::sqrt(vx * vx + vy * vy);
  return (vt >= kTPCActiveRMin && vt <= kTPCActiveRMax);
}
 
void GenTPCLoopers::setGeomProtection(bool protect)
{
  mGeomProtection = protect;
  if (mGeomProtection) {
    LOG(debug) << "TPC loopers geometrical protection: ON (accepting vertices with "
               << kTPCActiveRMin << " <= Vt <= " << kTPCActiveRMax << " cm)";
  } else {
    LOG(warning) << "TPC loopers geometrical protection: OFF - loopers will be generated outside the TPC active volume as well.";
  }
}
 
bool GenTPCLoopers::ownsCollision(const std::vector<o2::steer::EventPart>& parts, int source)
{
  // A collision is owned by the source with the highest ID among the event parts it holds, so that
  // exactly one of the simulations feeding a collision context (i.e. bkg, sgn) generates its loopers.
  // The signal comes after the background and has one event per collision, so it tiles the time
  // exactly; a reused background event can only cover the average of its gaps. This assumes that
  // signal events are not reused. The rule only looks at the collision itself, because the
  // background reads the global context and each signal its per-timeframe one.
  int highest = std::numeric_limits<int>::min();
  for (const auto& part : parts) {
    highest = std::max(highest, part.sourceID);
  }
  return highest == source;
}
 
bool GenTPCLoopers::ownsContextCollisions()
{
  // Loopers are a detector background which is continuous in time: event by event this generator
  // covers the time interval between consecutive collisions of the collision context, so a single
  // simulation tiles the whole timeframe exactly once. When several transport simulations feed the
  // same context (embedding), each of them would tile that same time axis again and the loopers
  // would end up in the digits as many times as there are simulations.
  // Every collision is therefore owned by exactly one source, the one with the highest source ID
  // among the event parts present, and only the owner generates loopers for it. Embedding without
  // replacement ("@" patterns) gives the signal the collisions it is embedded into and the
  // background the others; when the signal replaces the background ("r" patterns) each source
  // owns the collisions it is present in alone.
  const auto& simConfig = o2::conf::SimConfig::Instance();
  const auto contextAndPrefix = simConfig.getCollContextFilenameAndEventPrefix();
  const std::string contextName = contextAndPrefix.first.empty() ? "collisioncontext.root" : contextAndPrefix.first;
  const std::string& simPrefix = contextAndPrefix.second;
  if (!std::filesystem::exists(contextName)) {
    return true; // no context to share with anybody
  }
  std::unique_ptr<o2::steer::DigitizationContext> context(o2::steer::DigitizationContext::loadFromFile(contextName));
  if (!context) {
    LOG(warn) << "Could not read the collision context from " << contextName
              << "; loopers will be generated without checking for other sources";
    return true;
  }
  const auto source = context->findSimPrefix(simPrefix);
  if (source < 0) {
    if (context->getSimPrefixes().size() <= 1) {
      return true; // a single source: nobody to share the loopers with
    }
    LOG(warn) << "Simulation prefix " << simPrefix << " not found in the collision context " << contextName
              << ", which is shared by " << context->getSimPrefixes().size() << " sources; loopers will remain OFF";
    return false;
  }
  // count the collisions this source takes part in and those it owns
  unsigned int nPresent = 0, nOwned = 0;
  for (const auto& parts : context->getEventParts()) {
    if (ownsCollision(parts, source)) {
      nOwned++;
    }
    if (std::find_if(parts.begin(), parts.end(), [source](const auto& part) { return part.sourceID == source; }) != parts.end()) {
      nPresent++;
    }
  }
  LOG(info) << "Loopers: source " << simPrefix << " owns " << nOwned << " of the " << nPresent
            << " collisions it takes part in, out of " << context->getEventParts().size() << " in " << contextName;
  return nOwned > 0;
}
 
void GenTPCLoopers::setupContextCoverage()
{
  // Work out how much of the time axis every event of this simulation has to cover with loopers.
  // A collision stands for the time from itself to the next one (the last one to the end of the
  // frame), and an event covers the time of the collisions it is used at: when it is used at
  // several of them, as the background of an embedding production is, the loopers it carries are
  // digitised once per use, so it covers their average. Collisions owned by another source count
  // as zero. The loopers reaching the digits then add up to the reference number per orbit times
  // the time this source owns, however its events are spread over the collision context.
  mEventLooperTime.clear();
  mSourceID = mCollisionContext ? mCollisionContext->findSimPrefix(o2::conf::SimConfig::Instance().getCollContextFilenameAndEventPrefix().second) : -1;
  if (mSourceID < 0) {
    return; // no context information: every event keeps the interval of the collision of same index
  }
  const auto& parts = mCollisionContext->getEventParts();
  std::vector<int> nUses;
  double ownedTime = 0.;
  for (size_t c = 0; c < parts.size() && c < mInteractionTimeRecords.size(); ++c) {
    const double gap = (c + 1 < mInteractionTimeRecords.size())
                         ? mInteractionTimeRecords[c + 1].bc2ns() - mInteractionTimeRecords[c].bc2ns()
                         : mTimeEnd - mInteractionTimeRecords[c].bc2ns();
    const double owned = ownsCollision(parts[c], mSourceID) ? gap : 0.;
    ownedTime += owned;
    for (const auto& part : parts[c]) {
      if (part.sourceID != mSourceID) {
        continue;
      }
      if (part.entryID >= (int)mEventLooperTime.size()) { // entryIDs are dense, see SimConfig
        mEventLooperTime.resize(part.entryID + 1, 0.);
        nUses.resize(part.entryID + 1, 0);
      }
      mEventLooperTime[part.entryID] += owned;
      nUses[part.entryID]++;
    }
  }
  for (size_t e = 0; e < mEventLooperTime.size(); ++e) {
    mEventLooperTime[e] /= std::max(nUses[e], 1);
  }
  LOG(info) << "Loopers: the " << mEventLooperTime.size() << " events of this simulation share the "
            << ownedTime << " ns of the timeframe it owns";
}
 
GenTPCLoopers::GenTPCLoopers(std::string model_pairs, std::string model_compton,
                             std::string poisson, std::string gauss, std::string scaler_pair,
                             std::string scaler_compton)
{
  // Checking if the model files exist and are not empty
  std::ifstream model_file[2];
  model_file[0].open(model_pairs);
  model_file[1].open(model_compton);
  if (!model_file[0].is_open() || model_file[0].peek() == std::ifstream::traits_type::eof()) {
    LOG(fatal) << "Error: Pairs model file is empty or does not exist!";
    exit(1);
  }
  if (!model_file[1].is_open() || model_file[1].peek() == std::ifstream::traits_type::eof()) {
    LOG(fatal) << "Error: Compton model file is empty or does not exist!";
    exit(1);
  }
  model_file[0].close();
  model_file[1].close();
  // Checking if the scaler files exist and are not empty
  std::ifstream scaler_file[2];
  scaler_file[0].open(scaler_pair);
  scaler_file[1].open(scaler_compton);
  if (!scaler_file[0].is_open() || scaler_file[0].peek() == std::ifstream::traits_type::eof()) {
    LOG(fatal) << "Error: Pairs scaler file is empty or does not exist!";
    exit(1);
  }
  if (!scaler_file[1].is_open() || scaler_file[1].peek() == std::ifstream::traits_type::eof()) {
    LOG(fatal) << "Error: Compton scaler file is empty or does not exist!";
    exit(1);
  }
  scaler_file[0].close();
  scaler_file[1].close();
  // Checking if the poisson file exists and it's not empty
  if (poisson != "" && poisson != "None" && poisson != "none") {
    std::ifstream poisson_file(poisson);
    if (!poisson_file.is_open() || poisson_file.peek() == std::ifstream::traits_type::eof()) {
      LOG(fatal) << "Error: Poisson file is empty or does not exist!";
      exit(1);
    } else {
      poisson_file >> mPoisson[0] >> mPoisson[1] >> mPoisson[2];
      poisson_file.close();
      mPoissonSet = true;
    }
  }
  // Checking if the gauss file exists and it's not empty
  if (gauss != "" && gauss != "None" && gauss != "none") {
    std::ifstream gauss_file(gauss);
    if (!gauss_file.is_open() || gauss_file.peek() == std::ifstream::traits_type::eof()) {
      LOG(fatal) << "Error: Gauss file is empty or does not exist!";
      exit(1);
    } else {
      gauss_file >> mGauss[0] >> mGauss[1] >> mGauss[2] >> mGauss[3];
      gauss_file.close();
      mGaussSet = true;
    }
  }
  mONNX_pair = std::make_unique<ONNXGenerator>(global_env, model_pairs);
  mScaler_pair = std::make_unique<Scaler>();
  mScaler_pair->load(scaler_pair);
  mONNX_compton = std::make_unique<ONNXGenerator>(global_env, model_compton);
  mScaler_compton = std::make_unique<Scaler>();
  mScaler_compton->load(scaler_compton);
}
 
Bool_t GenTPCLoopers::generateEvent()
{
  // Clear the vector of pairs
  mGenPairs.clear();
  // Clear the vector of compton electrons
  mGenElectrons.clear();
  if (mFlatGas) {
    unsigned int nLoopers, nLoopersPairs, nLoopersCompton;
    LOG(debug) << "mCurrentEvent is " << mCurrentEvent;
    // the share of the timeframe this event has to cover, or, with no collision context to look
    // at, the interval between the collision of same index and the next one
    mTimeLimit = (mCurrentEvent < (int)mEventLooperTime.size())
                   ? mEventLooperTime[mCurrentEvent]
                   : ((mCurrentEvent < mInteractionTimeRecords.size() - 1) ? mInteractionTimeRecords[mCurrentEvent + 1].bc2ns() - mInteractionTimeRecords[mCurrentEvent].bc2ns() : mTimeEnd - mInteractionTimeRecords[mCurrentEvent].bc2ns());
    LOG(debug) << "Current event time: " << mTimeLimit << " ns";
    // With flat gas the number of loopers are adapted based on time interval widths
    // The denominator is either the LHC orbit (if mFlatGasOrbit is true) or the mean interaction time record interval
    nLoopers = mFlatGasOrbit ? (mFlatGasNumber * (mTimeLimit / o2::constants::lhc::LHCOrbitNS)) : (mFlatGasNumber * (mTimeLimit / mIntTimeRecMean));
    nLoopersPairs = static_cast<unsigned int>(std::round(nLoopers * mLoopsFractionPairs));
    nLoopersCompton = nLoopers - nLoopersPairs;
    SetNLoopers(nLoopersPairs, nLoopersCompton);
    LOG(info) << "Flat gas loopers: " << nLoopers << " (pairs: " << nLoopersPairs << ", compton: " << nLoopersCompton << ")";
    generateEvent(mTimeLimit);
    mCurrentEvent++;
  } else {
    // Set number of loopers if poissonian params are available
    if (mPoissonSet) {
      mNLoopersPairs = static_cast<unsigned int>(std::round(mMultiplier[0] * PoissonPairs()));
      LOG(debug) << "Generated loopers pairs (Poisson): " << mNLoopersPairs;
    }
    if (mGaussSet) {
      mNLoopersCompton = static_cast<unsigned int>(std::round(mMultiplier[1] * GaussianElectrons()));
      LOG(debug) << "Generated compton electrons (Gauss): " << mNLoopersCompton;
    }
    // Generate pairs
    for (int i = 0; i < mNLoopersPairs; ++i) {
      std::vector<double> pair = mONNX_pair->generate_sample();
      // Apply the inverse transformation using the scaler
      std::vector<double> transformed_pair = mScaler_pair->inverse_transform(pair);
      mGenPairs.push_back(transformed_pair);
    }
    // Generate compton electrons
    for (int i = 0; i < mNLoopersCompton; ++i) {
      std::vector<double> electron = mONNX_compton->generate_sample();
      // Apply the inverse transformation using the scaler
      std::vector<double> transformed_electron = mScaler_compton->inverse_transform(electron);
      mGenElectrons.push_back(transformed_electron);
    }
  }
  return true;
}
 
Bool_t GenTPCLoopers::generateEvent(double time_limit)
{
  LOG(info) << "Time constraint for loopers: " << time_limit << " ns";
  // Generate pairs
  for (int i = 0; i < mNLoopersPairs; ++i) {
    std::vector<double> pair = mONNX_pair->generate_sample();
    // Apply the inverse transformation using the scaler
    std::vector<double> transformed_pair = mScaler_pair->inverse_transform(pair);
    transformed_pair[9] = gRandom->Uniform(0., time_limit); // Regenerate time, scaling is not needed because time_limit is already in nanoseconds
    mGenPairs.push_back(transformed_pair);
  }
  // Generate compton electrons
  for (int i = 0; i < mNLoopersCompton; ++i) {
    std::vector<double> electron = mONNX_compton->generate_sample();
    // Apply the inverse transformation using the scaler
    std::vector<double> transformed_electron = mScaler_compton->inverse_transform(electron);
    transformed_electron[6] = gRandom->Uniform(0., time_limit); // Regenerate time, scaling is not needed because time_limit is already in nanoseconds
    mGenElectrons.push_back(transformed_electron);
  }
  LOG(info) << "Generated Particles with time limit";
  return true;
}
 
std::vector<TParticle> GenTPCLoopers::importParticles()
{
  std::vector<TParticle> particles;
  const double mass_e = TDatabasePDG::Instance()->GetParticle(11)->Mass();
  const double mass_p = TDatabasePDG::Instance()->GetParticle(-11)->Mass();
  mNSkippedPairs = 0;
  mNSkippedCompton = 0;
  // Get looper pairs from the event
  for (auto& pair : mGenPairs) {
    double px_e, py_e, pz_e, px_p, py_p, pz_p;
    double vx, vy, vz, time;
    double e_etot, p_etot;
    // The generative model is not currently fully constrained to the TPC geometry, so it places
    // significant fraction of the vertices outside the drift gas.
    // These are now dropped before they reach the transport.
    if (mGeomProtection && !isInTPCActiveVolume(pair[6], pair[7])) {
      mNSkippedPairs++;
      continue;
    }
    px_e = pair[0];
    py_e = pair[1];
    pz_e = pair[2];
    px_p = pair[3];
    py_p = pair[4];
    pz_p = pair[5];
    vx = pair[6];
    vy = pair[7];
    vz = pair[8];
    time = pair[9];
    e_etot = TMath::Sqrt(px_e * px_e + py_e * py_e + pz_e * pz_e + mass_e * mass_e);
    p_etot = TMath::Sqrt(px_p * px_p + py_p * py_p + pz_p * pz_p + mass_p * mass_p);
    // Push the electron
    TParticle electron(11, 1, -1, -1, -1, -1, px_e, py_e, pz_e, e_etot, vx, vy, vz, time / 1e9);
    // Setting HepMC status code != 1 to avoid detecting the loopers as physical primaries
    electron.SetStatusCode(o2::mcgenstatus::MCGenStatusEncoding(2, 0).fullEncoding);
    electron.SetBit(ParticleStatus::kToBeDone, true);
    particles.push_back(electron);
    // Push the positron
    TParticle positron(-11, 1, -1, -1, -1, -1, px_p, py_p, pz_p, p_etot, vx, vy, vz, time / 1e9);
    positron.SetStatusCode(o2::mcgenstatus::MCGenStatusEncoding(2, 0).fullEncoding);
    positron.SetBit(ParticleStatus::kToBeDone, true);
    particles.push_back(positron);
  }
  // Get compton electrons from the event
  for (auto& compton : mGenElectrons) {
    double px, py, pz;
    double vx, vy, vz, time;
    double etot;
    if (mGeomProtection && !isInTPCActiveVolume(compton[3], compton[4])) {
      mNSkippedCompton++;
      continue;
    }
    px = compton[0];
    py = compton[1];
    pz = compton[2];
    vx = compton[3];
    vy = compton[4];
    vz = compton[5];
    time = compton[6];
    etot = TMath::Sqrt(px * px + py * py + pz * pz + mass_e * mass_e);
    // Push the electron
    TParticle electron(11, 1, -1, -1, -1, -1, px, py, pz, etot, vx, vy, vz, time / 1e9);
    // Setting HepMC status code != 1 to avoid detecting the loopers as physical primaries
    electron.SetStatusCode(o2::mcgenstatus::MCGenStatusEncoding(2, 0).fullEncoding);
    electron.SetBit(ParticleStatus::kToBeDone, true);
    particles.push_back(electron);
  }
 
  return particles;
}
 
unsigned int GenTPCLoopers::PoissonPairs()
{
  unsigned int poissonValue;
  do {
    // Generate a Poisson-distributed random number with mean mPoisson[0]
    poissonValue = mRandGen.Poisson(mPoisson[0]);
  } while (poissonValue < mPoisson[1] || poissonValue > mPoisson[2]); // Regenerate if out of range
 
  return poissonValue;
}
 
unsigned int GenTPCLoopers::GaussianElectrons()
{
  unsigned int gaussValue;
  do {
    // Generate a Normal-distributed random number with mean mGass[0] and stddev mGauss[1]
    gaussValue = mRandGen.Gaus(mGauss[0], mGauss[1]);
  } while (gaussValue < mGauss[2] || gaussValue > mGauss[3]); // Regenerate if out of range
 
  return gaussValue;
}
 
void GenTPCLoopers::SetNLoopers(unsigned int nsig_pair, unsigned int nsig_compton)
{
  if (mFlatGas) {
    mNLoopersPairs = nsig_pair;
    mNLoopersCompton = nsig_compton;
  } else {
    if (mPoissonSet) {
      LOG(info) << "Poissonian parameters correctly loaded.";
    } else {
      mNLoopersPairs = nsig_pair;
    }
    if (mGaussSet) {
      LOG(info) << "Gaussian parameters correctly loaded.";
    } else {
      mNLoopersCompton = nsig_compton;
    }
  }
}
 
void GenTPCLoopers::SetMultiplier(const std::array<float, 2>& mult)
{
  // Multipliers will work only if the poissonian and gaussian parameters are set
  // otherwise they will be ignored
  if (mult[0] < 0 || mult[1] < 0) {
    LOG(fatal) << "Error: Multiplier values must be non-negative!";
    exit(1);
  } else {
    LOG(info) << "Multiplier values set to: Pair = " << mult[0] << ", Compton = " << mult[1];
    mMultiplier[0] = mult[0];
    mMultiplier[1] = mult[1];
  }
}
 
void GenTPCLoopers::setFlatGas(Bool_t flat, Int_t number, Int_t nloopers_orbit)
{
  mFlatGas = flat;
  if (mFlatGas) {
    if (nloopers_orbit > 0) {
      mFlatGasOrbit = true;
      mFlatGasNumber = nloopers_orbit;
      LOG(info) << "Flat gas loopers will be generated using orbit reference.";
    } else {
      mFlatGasOrbit = false;
      if (number < 0) {
        LOG(warn) << "Warning: Number of loopers per event must be non-negative! Switching option off.";
        mFlatGas = false;
        mFlatGasNumber = -1;
      } else {
        mFlatGasNumber = number;
      }
    }
    if (mFlatGas) {
      mContextFile = std::filesystem::exists("collisioncontext.root") ? TFile::Open("collisioncontext.root") : nullptr;
      mCollisionContext = mContextFile ? (o2::steer::DigitizationContext*)mContextFile->Get("DigitizationContext") : nullptr;
      mInteractionTimeRecords = mCollisionContext ? mCollisionContext->getEventRecords() : std::vector<o2::InteractionTimeRecord>{};
      const auto& hbfUtils = o2::raw::HBFUtils::Instance();
      if (mInteractionTimeRecords.empty()) {
        // A timeframe can legitimately contain no collision at all when the interaction rate is
        // low. No event is transported in that case, so nothing below is ever used; take the
        // extent of the timeframe from HBFUtils rather than from the (absent) collisions.
        LOG(warn) << "No interaction time records in the collision context; this timeframe holds no collision";
        o2::InteractionRecord tfEndIR(0, hbfUtils.orbitFirstSampled + hbfUtils.nHBFPerTF);
        mTimeEnd = tfEndIR.bc2ns();
      } else {
        LOG(info) << "Interaction Time records has " << mInteractionTimeRecords.size() << " entries.";
        mCollisionContext->printCollisionSummary();
        for (int c = 0; c < (int)mInteractionTimeRecords.size() - 1; c++) {
          mIntTimeRecMean += mInteractionTimeRecords[c + 1].bc2ns() - mInteractionTimeRecords[c].bc2ns();
        }
        if (mInteractionTimeRecords.size() > 1) {
          mIntTimeRecMean /= (mInteractionTimeRecords.size() - 1); // Average interaction time record used as reference
        } else {
          // a single collision gives no spacing to average; use the one implied by the rate
          auto rate = mCollisionContext->getDigitizerInteractionRate();
          mIntTimeRecMean = rate > 0. ? 1.e9 / rate : (double)o2::constants::lhc::LHCOrbitNS;
          LOG(info) << "Only one collision in this timeframe; taking " << mIntTimeRecMean
                    << " ns as the mean interaction spacing from the interaction rate";
        }
        // Get the start time of the second orbit after the last interaction record
        const auto& lastIR = mInteractionTimeRecords.back();
        o2::InteractionRecord finalOrbitIR(0, lastIR.orbit + 2); // Final orbit, BC = 0
        mTimeEnd = finalOrbitIR.bc2ns();
        LOG(debug) << "Final orbit start time: " << mTimeEnd << " ns while last interaction record time is " << mInteractionTimeRecords.back().bc2ns() << " ns";
        setupContextCoverage();
      }
    }
  } else {
    mFlatGasNumber = -1;
  }
  LOG(info) << "Flat gas loopers: " << (mFlatGas ? "ON" : "OFF") << ", Reference loopers number per " << (mFlatGasOrbit ? "orbit " : "event ") << mFlatGasNumber;
}
 
void GenTPCLoopers::setFractionPairs(float fractionPairs)
{
  if (fractionPairs < 0 || fractionPairs > 1) {
    LOG(fatal) << "Error: Loops fraction for pairs must be in the range [0, 1].";
    exit(1);
  }
  mLoopsFractionPairs = fractionPairs;
  LOG(info) << "Pairs fraction set to: " << mLoopsFractionPairs;
}
 
void GenTPCLoopers::SetRate(const std::string& rateFile, bool isPbPb = true, int intRate)
{
  // Checking if the rate file exists and is not empty
  TFile rate_file(rateFile.c_str(), "READ");
  if (!rate_file.IsOpen() || rate_file.IsZombie()) {
    LOG(fatal) << "Error: Rate file is empty or does not exist!";
    exit(1);
  }
  const char* fitName = isPbPb ? "fitPbPb" : "fitpp";
  auto fit = (TF1*)rate_file.Get(fitName);
  if (!fit) {
    LOG(fatal) << "Error: Could not find fit function '" << fitName << "' in rate file!";
    exit(1);
  }
  mInteractionRate = intRate;
  if (mInteractionRate < 0) {
    mContextFile = std::filesystem::exists("collisioncontext.root") ? TFile::Open("collisioncontext.root") : nullptr;
    if (!mContextFile || mContextFile->IsZombie()) {
      LOG(fatal) << "Error: Interaction rate not provided and collision context file not found!";
      exit(1);
    }
    mCollisionContext = (o2::steer::DigitizationContext*)mContextFile->Get("DigitizationContext");
    mInteractionRate = std::floor(mCollisionContext->getDigitizerInteractionRate());
    LOG(info) << "Interaction rate retrieved from collision context: " << mInteractionRate << " Hz";
    if (mInteractionRate < 0) {
      LOG(fatal) << "Error: Invalid interaction rate retrieved from collision context!";
      exit(1);
    }
  }
  auto ref = static_cast<int>(std::floor(fit->Eval(mInteractionRate / 1000.))); // fit expects rate in kHz
  rate_file.Close();
  if (ref <= 0) {
    LOG(fatal) << "Computed flat gas number reference per orbit is <=0";
    exit(1);
  } else {
    LOG(info) << "Set flat gas number to " << ref << " loopers per orbit using " << fitName << " from " << mInteractionRate << " Hz interaction rate.";
    auto flat = true;
    setFlatGas(flat, -1, ref);
  }
}
 
void GenTPCLoopers::SetAdjust(float adjust)
{
  if (mFlatGas && mFlatGasOrbit && adjust >= -1.f && adjust != 0.f) {
    LOG(info) << "Adjusting flat gas number per orbit by " << adjust * 100.f << "%";
    mFlatGasNumber = static_cast<int>(std::round(mFlatGasNumber * (1.f + adjust)));
    LOG(info) << "New flat gas number per orbit: " << mFlatGasNumber;
  }
}
 
} // namespace eventgen
} // namespace o2