Introducing non-uniform mu InteractionSampler

Provides a novel InteractionSampler for collision structure
which is able to sample (orbit,bc) values according to an
externally given hNBcVTX distribution (obtained from FIT, EventSelection)

A unit test which shows and tests the new feature.

Fixes https://its.cern.ch/jira/browse/O2-6450

Author: sawenzel

DataFormats/simulation/CMakeLists.txt

@@ -55,6 +55,11 @@ o2_target_root_dictionary(
 # * src/SimulationDataLinkDef.h
 # * and not src/SimulationDataFormatLinkDef.h
 
+o2_add_test(InteractionSampler
+            SOURCES test/testInteractionSampler.cxx
+            COMPONENT_NAME SimulationDataFormat
+            PUBLIC_LINK_LIBRARIES O2::SimulationDataFormat)
+
 o2_add_test(BasicHits
             SOURCES test/testBasicHits.cxx
             COMPONENT_NAME SimulationDataFormat

DataFormats/simulation/include/SimulationDataFormat/InteractionSampler.h

@@ -22,6 +22,7 @@
 #include "CommonDataFormat/BunchFilling.h"
 #include "CommonConstants/LHCConstants.h"
 #include "MathUtils/RandomRing.h"
+#include <TH1F.h>
 
 namespace o2
 {
@@ -130,6 +131,30 @@ class FixedSkipBC_InteractionSampler : public InteractionSampler
   ClassDef(FixedSkipBC_InteractionSampler, 1);
 };
 
+// A version of the interaction sampler which can sample according to non-uniform mu(bc) as
+// observed during data taking.
+class NonUniformMuInteractionSampler : public InteractionSampler
+{
+ public:
+  NonUniformMuInteractionSampler() : InteractionSampler() { mBCIntensityScales.resize(o2::constants::lhc::LHCMaxBunches, 1); }
+  bool setBCIntensityScales(const std::vector<float>& scales_from_vector);
+  bool setBCIntensityScales(const TH1F& scales_from_histo); // initialize scales
+
+  // helper function to determine the scales from a histogram (count from event selection analysis)
+  std::vector<float> determineBCIntensityScalesFromHistogram(const TH1F& scales_from_histo);
+
+  const std::vector<float>& getBCIntensityScales() const { return mBCIntensityScales; }
+
+ protected:
+  int simulateInteractingBC() override;
+  int getBCJump() const;
+
+ private:
+  // non-uniformity
+  std::vector<float> mBCIntensityScales;
+  ClassDef(NonUniformMuInteractionSampler, 1);
+};
+
 } // namespace steer
 } // namespace o2
 

DataFormats/simulation/src/InteractionSampler.cxx

@@ -115,8 +115,8 @@ const o2::InteractionTimeRecord& InteractionSampler::generateCollisionTime()
 int InteractionSampler::simulateInteractingBC()
 {
   // Returns number of collisions assigned to selected BC
-
   nextCollidingBC(mBCJumpGenerator.getNextValue());
+
   // once BC is decided, enforce at least one interaction
   int ncoll = mNCollBCGenerator.getNextValue();
 
@@ -162,3 +162,97 @@ void InteractionSampler::setBunchFilling(const std::string& bcFillingFile)
   mBCFilling = *bc;
   delete bc;
 }
+
+// ________________________________________________
+bool NonUniformMuInteractionSampler::setBCIntensityScales(const std::vector<float>& scales_from_vector)
+{
+  // Sets the intensity scales per bunch crossing index
+  // The length of this vector needs to be compatible with the bunch filling chosen
+  mBCIntensityScales = scales_from_vector;
+
+  if (scales_from_vector.size() != mInteractingBCs.size()) {
+    LOG(error) << "Scaling factors and bunch filling scheme are not compatible. Not doing anything";
+    return false;
+  }
+
+  // since this is a scaling factor, we need to normalize the max to 1
+  float max_val = 0.;
+  for (auto v : mBCIntensityScales) {
+    max_val = std::max(max_val, v);
+  }
+  // normalize
+  for (auto& v : mBCIntensityScales) {
+    if (v < 0.) {
+      v = std::abs(v);
+    } else {
+      v /= max_val;
+    }
+  }
+  return true;
+}
+
+// ________________________________________________
+
+bool NonUniformMuInteractionSampler::setBCIntensityScales(const TH1F& hist)
+{
+  return setBCIntensityScales(determineBCIntensityScalesFromHistogram(hist));
+}
+
+std::vector<float> NonUniformMuInteractionSampler::determineBCIntensityScalesFromHistogram(const TH1F& hist)
+{
+  std::vector<float> scales;
+  // we go through the BCs and query the count from histogram
+  for (auto bc : mInteractingBCs) {
+    scales.push_back(hist.GetBinContent(bc + 1));
+  }
+  return scales;
+}
+
+int NonUniformMuInteractionSampler::getBCJump() const
+{
+  auto muFunc = [this](int bc_position) {
+    return mBCIntensityScales[bc_position % mInteractingBCs.size()] * mMuBC;
+  };
+
+  double U = gRandom->Rndm();    // uniform (0,1)
+  double T = -std::log(1.0 - U); // threshold
+  double sumMu = 0.0;
+  int offset = 0;
+  auto bcStart = mCurrBCIdx; // the current bc
+  int maxSearch = o2::constants::lhc::LHCMaxBunches;
+
+  while (offset < maxSearch) {
+    auto mu_here = muFunc(bcStart + offset); // mu at next BC
+    sumMu += mu_here;
+    if (sumMu >= T) {
+      break; // found BC with at least one collision
+    }
+    ++offset;
+  }
+  return offset;
+}
+
+int NonUniformMuInteractionSampler::simulateInteractingBC()
+{
+  nextCollidingBC(getBCJump());
+
+  auto muFunc = [this](int bc_position) {
+    return mBCIntensityScales[bc_position % mInteractingBCs.size()] * mMuBC;
+  };
+
+  // now sample number of collisions in chosenBC, conditioned >=1:
+  double mu_chosen = muFunc(mCurrBCIdx); // or does it need to be mCurrBCIdx
+  int ncoll = 0;
+  do {
+    ncoll = gRandom->Poisson(mu_chosen);
+  } while (ncoll == 0);
+
+  // assign random time withing a bunch
+  for (int i = ncoll; i--;) {
+    mTimeInBC.push_back(mCollTimeGenerator.getNextValue());
+  }
+  if (ncoll > 1) { // sort in DECREASING time order (we are reading vector from the end)
+    std::sort(mTimeInBC.begin(), mTimeInBC.end(), [](const float a, const float b) { return a > b; });
+  }
+  return ncoll;
+}

DataFormats/simulation/src/SimulationDataLinkDef.h

@@ -25,6 +25,7 @@
 
 #pragma link C++ class o2::steer::InteractionSampler + ;
 #pragma link C++ class o2::steer::FixedSkipBC_InteractionSampler + ;
+#pragma link C++ class o2::steer::NonUniformMuInteractionSampler + ;
 #pragma link C++ class o2::sim::StackParam + ;
 #pragma link C++ class o2::conf::ConfigurableParamHelper < o2::sim::StackParam> + ;
 #pragma link C++ class o2::MCTrackT < double> + ;