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> + ;

DataFormats/simulation/test/testInteractionSampler.cxx

@@ -0,0 +1,76 @@
+// Copyright 2019-2020 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.
+
+#define BOOST_TEST_MODULE Test InteractionSampler class
+#define BOOST_TEST_MAIN
+#define BOOST_TEST_DYN_LINK
+
+#include <boost/test/unit_test.hpp>
+#include "SimulationDataFormat/InteractionSampler.h"
+#include "CCDB/BasicCCDBManager.h"
+#include "DataFormatsParameters/AggregatedRunInfo.h"
+#include "DataFormatsParameters/GRPLHCIFData.h"
+#include "TFile.h"
+#include "TGrid.h"
+#include <TH1F.h>
+
+namespace o2
+{
+
+BOOST_AUTO_TEST_CASE(NonUniformSampler)
+{
+  auto run_number = 559827;
+  auto runInfo = o2::parameters::AggregatedRunInfo::buildAggregatedRunInfo(o2::ccdb::BasicCCDBManager::instance(), run_number);
+
+  o2::steer::NonUniformMuInteractionSampler sampler;
+  sampler.setBunchFilling(runInfo.grpLHC->getBunchFilling());
+
+  TGrid::Connect("alien");
+  if (gGrid) {
+    // the test distribution provided by Igor Altsybeev
+    auto distr_file = TFile::Open("alien:///alice/cern.ch/user/s/swenzel/AliceO2_TestData/NBcVTX_559827/hBcTVX_data_PbPb_24ar_559827.root");
+
+    //
+    if (distr_file && !distr_file->IsZombie()) {
+      auto hist = distr_file->Get<TH1F>("hBcTVX");
+      if (hist) {
+        sampler.init();
+        sampler.setBCIntensityScales(*hist);
+
+        // sample into a vector of a certain size
+        std::vector<o2::InteractionTimeRecord> samples;
+
+        int N = 100000;
+        samples.resize(N);
+
+        sampler.generateCollisionTimes(samples);
+
+        // fill an output histogram
+        auto output_hist = (TH1F*)hist->Clone("h2"); // make a full copy
+        output_hist->Reset();
+
+        for (const auto& sample : samples) {
+          output_hist->Fill(sample.bc);
+        }
+
+        // Write out
+        auto fout = TFile::Open("NBCVTX_out.root", "RECREATE");
+        fout->WriteObject(output_hist, "NBcVTX");
+        fout->Close();
+
+        // compare mean values of original and newly sampled hist
+        BOOST_CHECK_CLOSE(hist->GetMean(), output_hist->GetMean(), 0.5);
+      }
+    }
+  }
+}
+
+} // namespace o2