Integrate non-uniform InteractionSampler into CollisionContextTool

Author: sawenzel

DataFormats/simulation/src/InteractionSampler.cxx

@@ -202,6 +202,9 @@ bool NonUniformMuInteractionSampler::setBCIntensityScales(const TH1F& hist)
 
 std::vector<float> NonUniformMuInteractionSampler::determineBCIntensityScalesFromHistogram(const TH1F& hist)
 {
+  if (mInteractingBCs.size() == 0) {
+    LOG(error) << " Initialize bunch crossing scheme before assigning scales";
+  }
   std::vector<float> scales;
   // we go through the BCs and query the count from histogram
   for (auto bc : mInteractingBCs) {

Steer/src/CollisionContextTool.cxx

@@ -64,6 +64,8 @@ struct Options {
                                   // This is useful when someone else is creating the contexts (MC-data embedding) and we
                                   // merely want to pass these through. If this is given, we simply take the timeframe ID, number of orbits
                                   // and copy the right amount of timeframes into the destination folder (implies individualTFextraction)
+  std::string nontrivial_mu_distribution = ""; // path to fetch a non-uniform MC(BC) distribution for the interaction sampler
+                                               // can be: (a) ccdb, (b) a ROOT file with the histogram included
 };
 
 enum class InteractionLockMode {
@@ -216,8 +218,8 @@ bool parseOptions(int argc, char* argv[], Options& optvalues)
     "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]")(
-    "import-external", bpo::value<std::string>(&optvalues.external_path)->default_value(""),
-    "Take collision contexts (per timeframe) from external files for instance for data-anchoring use-case. Needs timeframeID and number of orbits to be given as well.");
+    "import-external", bpo::value<std::string>(&optvalues.external_path)->default_value(""), "Take collision contexts (per timeframe) from external files for instance for data-anchoring use-case. Needs timeframeID and number of orbits to be given as well.")(
+    "nontrivial-mu-distribution", bpo::value<std::string>(&optvalues.nontrivial_mu_distribution)->default_value(""), "Distribution for MU(BC)");
 
   options.add_options()("help,h", "Produce help message.");
 
@@ -397,6 +399,37 @@ int main(int argc, char* argv[])
     auto mode = ispecs[id].syncmode;
     if (mode == InteractionLockMode::NOLOCK) {
       auto sampler = std::make_unique<o2::steer::InteractionSampler>();
+      TH1F* mu_hist = nullptr;
+
+      // we check if there is a realistic bunch crossing distribution available
+      const auto& mu_distr_source = options.nontrivial_mu_distribution;
+      if (mu_distr_source.size() > 0) {
+        if (mu_distr_source == "ccdb") {
+          auto& inst = o2::ccdb::BasicCCDBManager::instance();
+          // we fetch it from CCDB ... but allow empty response
+          auto backup = inst.getFatalWhenNull();
+          inst.setFatalWhenNull(false);
+          auto local_hist = inst.getForTimeStamp<TH1F>("Users/s/swenzel/GLO/HBCTVX", options.timestamp);
+          inst.setFatalWhenNull(backup);
+          mu_hist = (TH1F*)(local_hist->Clone("h2")); // we need to clone since ownership of local_hist is with TFile
+        } else {
+          // we interpret the file as a ROOT file and open it to extract the wanted histogram
+          auto mudistr_file = TFile::Open(mu_distr_source.c_str(), "OPEN");
+          if (mudistr_file && !mudistr_file->IsZombie()) {
+            auto local_hist = mudistr_file->Get<TH1F>("hBcTVX");
+            mu_hist = (TH1F*)(local_hist->Clone("h2")); // we need to clone since ownership of local_hist is with TFile
+            mudistr_file->Close();
+          }
+        }
+        if (mu_hist) {
+          LOG(info) << "Found an external mu distribution with mean BC value " << mu_hist->GetMean();
+
+          // do some checks
+
+          // reset to correct interaction Sampler type
+          sampler.reset(new o2::steer::NonUniformMuInteractionSampler());
+        }
+      }
 
       // for debug purposes: allows to instantiate trivial sampler
       if (const char* env = getenv("ALICEO2_ENFORCE_TRIVIAL_BC_SAMPLER")) {
@@ -418,11 +451,17 @@ int main(int argc, char* argv[])
       if (!options.bcpatternfile.empty()) {
         setBCFillingHelper(*sampler, options.bcpatternfile);
       }
+      sampler->init();
+      if (auto sampler_cast = dynamic_cast<o2::steer::NonUniformMuInteractionSampler*>(sampler.get())) {
+        if (mu_hist) {
+          sampler_cast->setBCIntensityScales(*mu_hist);
+        }
+      }
+
       o2::InteractionTimeRecord record;
       // this loop makes sure that the first collision is within the range of orbits asked (if noEmptyTF is enabled)
       do {
         sampler->setFirstIR(o2::InteractionRecord(options.firstBC, orbitstart));
-        sampler->init();
         record = sampler->generateCollisionTime();
       } while (options.noEmptyTF && usetimeframelength && record.orbit >= orbitstart + orbits_total);
       int count = 0;