TPC: Adding streamer for FastTransform in reconstruction

Common/Utils/CMakeLists.txt

@@ -25,7 +25,7 @@ o2_add_library(CommonUtils
                        src/IRFrameSelector.cxx
                        src/DebugStreamer.cxx
                PUBLIC_LINK_LIBRARIES ROOT::Hist ROOT::Tree Boost::iostreams O2::CommonDataFormat O2::Headers
-                                     FairLogger::FairLogger O2::MathUtils)
+                                     FairLogger::FairLogger O2::MathUtils TBB::tbb)
 
 o2_target_root_dictionary(CommonUtils
                           HEADERS include/CommonUtils/TreeStream.h

Common/Utils/include/CommonUtils/DebugStreamer.h

@@ -21,6 +21,7 @@
 #include "CommonUtils/ConfigurableParamHelper.h"
 #if defined(DEBUG_STREAMER)
 #include "CommonUtils/TreeStreamRedirector.h"
+#include <tbb/concurrent_unordered_map.h>
 #endif
 #endif
 
@@ -29,10 +30,12 @@ namespace o2::utils
 
 /// struct defining the flags which can be used to check if a certain debug streamer is used
 enum StreamFlags {
-  streamdEdx = 1 << 0,         ///< stream corrections and cluster properties used for the dE/dx
-  streamDigitFolding = 1 << 1, ///< stream ion tail and saturatio information
-  streamDigits = 1 << 2,       ///< stream digit information
-  streamITCorr = 1 << 4,       ///< stream ion tail correction information
+  streamdEdx = 1 << 0,          ///< stream corrections and cluster properties used for the dE/dx
+  streamDigitFolding = 1 << 1,  ///< stream ion tail and saturatio information
+  streamDigits = 1 << 2,        ///< stream digit information
+  streamFastTransform = 1 << 3, ///< stream tpc fast transform
+  streamITCorr = 1 << 4,        ///< stream ion tail correction information
+  streamDistortionsSC = 1 << 5, ///< stream distortions applied in the TPC space-charge class (used for example in the tpc digitizer)
 };
 
 #if !defined(GPUCA_GPUCODE) && !defined(GPUCA_STANDALONE)
@@ -59,16 +62,33 @@ class DebugStreamer
 // CPU implementation of the class
 #if !defined(GPUCA_GPUCODE) && !defined(GPUCA_STANDALONE) && defined(DEBUG_STREAMER)
  public:
+  /// default constructor
+  DebugStreamer();
+
   /// set the streamer i.e. create the output file and set up the streamer
   /// \param outFile output file name without .root suffix
   /// \param option RECREATE or UPDATE
-  void setStreamer(const char* outFile, const char* option);
+  /// \param id unique id for given streamer (i.e. for defining unique streamers for each thread)
+  void setStreamer(const char* outFile, const char* option, const size_t id = getCPUID());
 
   /// \return returns if the streamer is set
-  bool isStreamerSet() const { return mTreeStreamer ? true : false; }
+  /// \param id unique id of streamer
+  bool isStreamerSet(const size_t id = getCPUID()) const { return getStreamerPtr(id); }
+
+  /// flush TTree for given ID to disc
+  /// \param id unique id of streamer
+  void flush(const size_t id);
+
+  /// flush all TTrees to disc
+  void flush();
+
+  /// \return returns streamer object for given id
+  /// \param id unique id of streamer
+  o2::utils::TreeStreamRedirector& getStreamer(const size_t id = getCPUID()) { return *(mTreeStreamer[id]); }
 
   /// \return returns streamer object
-  o2::utils::TreeStreamRedirector& getStreamer() { return *mTreeStreamer; }
+  /// \param id unique id of streamer
+  o2::utils::TreeStreamRedirector* getStreamerPtr(const size_t id = getCPUID()) const;
 
   /// \return returns streamer level i.e. what will be written to file
   static StreamFlags getStreamFlags() { return ParameterDebugStreamer::Instance().StreamLevel; }
@@ -77,11 +97,13 @@ class DebugStreamer
   static size_t getCPUID();
 
   /// \return returns number of trees in the streamer
-  int getNTrees() const;
+  /// \param id unique id of streamer
+  int getNTrees(const size_t id = getCPUID()) const;
 
   /// \return returns an unique branch name which is not already written in the file
   /// \param tree name of the tree for which to get a unique tree name
-  std::string getUniqueTreeName(const char* tree) const;
+  /// \param id unique id of streamer
+  std::string getUniqueTreeName(const char* tree, const size_t id = getCPUID()) const;
 
   /// set directly the debug level
   static void setStreamFlags(const StreamFlags streamFlags) { o2::conf::ConfigurableParam::setValue("DebugStreamerParam", "StreamLevel", static_cast<int>(streamFlags)); }
@@ -102,7 +124,9 @@ class DebugStreamer
   static void mergeTrees(const char* inpFile, const char* outFile, const char* option = "fast");
 
  private:
-  std::unique_ptr<o2::utils::TreeStreamRedirector> mTreeStreamer; ///< streamer which is used for the debugging
+  using StreamersPerFlag = tbb::concurrent_unordered_map<size_t, std::unique_ptr<o2::utils::TreeStreamRedirector>>;
+  StreamersPerFlag mTreeStreamer; ///< streamer which is used for the debugging
+
 #else
 
   // empty implementation of the class for GPU or when the debug streamer is not build for CPU
@@ -126,14 +150,16 @@ class DebugStreamer
     }
   };
 
-  GPUd() StreamerDummy getStreamer() const { return StreamerDummy{}; };
+  GPUd() StreamerDummy getStreamer(const int id = 0) const { return StreamerDummy{}; };
 
   template <typename Type>
-  GPUd() StreamerDummy getUniqueTreeName(Type) const
+  GPUd() StreamerDummy getUniqueTreeName(Type, const int id = 0) const
   {
     return StreamerDummy{};
   }
 
+  GPUd() void flush() const {};
+
 #endif
 };
 

Common/Utils/src/DebugStreamer.cxx

@@ -21,32 +21,65 @@ O2ParamImpl(o2::utils::ParameterDebugStreamer);
 
 #if !defined(GPUCA_GPUCODE) && !defined(GPUCA_STANDALONE) && defined(DEBUG_STREAMER)
 
-void o2::utils::DebugStreamer::setStreamer(const char* outFile, const char* option)
+o2::utils::DebugStreamer::DebugStreamer()
 {
-  if (!isStreamerSet()) {
-    ROOT::EnableThreadSafety();
-    mTreeStreamer = std::make_unique<o2::utils::TreeStreamRedirector>(fmt::format("{}_{}.root", outFile, getCPUID()).data(), option);
+  ROOT::EnableThreadSafety();
+}
+
+void o2::utils::DebugStreamer::setStreamer(const char* outFile, const char* option, const size_t id)
+{
+  if (!isStreamerSet(id)) {
+    mTreeStreamer[id] = std::make_unique<o2::utils::TreeStreamRedirector>(fmt::format("{}_{}.root", outFile, id).data(), option);
+  }
+}
+
+void o2::utils::DebugStreamer::flush(const size_t id)
+{
+  if (isStreamerSet(id)) {
+    mTreeStreamer[id].reset();
+  }
+}
+
+void o2::utils::DebugStreamer::flush()
+{
+  for (const auto& pair : mTreeStreamer) {
+    flush(pair.first);
   }
 }
 
-std::string o2::utils::DebugStreamer::getUniqueTreeName(const char* tree) const { return fmt::format("{}_{}", tree, getNTrees()); }
+std::string o2::utils::DebugStreamer::getUniqueTreeName(const char* tree, const size_t id) const { return fmt::format("{}_{}", tree, getNTrees(id)); }
 
 size_t o2::utils::DebugStreamer::getCPUID() { return std::hash<std::thread::id>{}(std::this_thread::get_id()); }
 
-int o2::utils::DebugStreamer::getNTrees() const { return mTreeStreamer->GetFile()->GetListOfKeys()->GetEntries(); }
+o2::utils::TreeStreamRedirector* o2::utils::DebugStreamer::getStreamerPtr(const size_t id) const
+{
+  auto it = mTreeStreamer.find(id);
+  if (it != mTreeStreamer.end()) {
+    return (it->second).get();
+  } else {
+    return nullptr;
+  }
+}
+
+int o2::utils::DebugStreamer::getNTrees(const size_t id) const { return isStreamerSet(id) ? getStreamerPtr(id)->GetFile()->GetListOfKeys()->GetEntries() : -1; }
 
 void o2::utils::DebugStreamer::mergeTrees(const char* inpFile, const char* outFile, const char* option)
 {
   TFile fInp(inpFile, "READ");
-  TList list;
+  std::unordered_map<int, TList> lists;
   for (TObject* keyAsObj : *fInp.GetListOfKeys()) {
     const auto key = dynamic_cast<TKey*>(keyAsObj);
-    list.Add((TTree*)fInp.Get(key->GetName()));
+    TTree* tree = (TTree*)fInp.Get(key->GetName());
+    // perform simple check on the number of entries to merge only TTree with same content (ToDo: Do check on name of branches)
+    const int entries = tree->GetListOfBranches()->GetEntries();
+    lists[entries].Add(tree);
   }
 
   TFile fOut(outFile, "RECREATE");
-  auto tree = TTree::MergeTrees(&list, option);
-  fOut.WriteObject(tree, "tree");
+  for (auto& list : lists) {
+    auto tree = TTree::MergeTrees(&list.second, option);
+    fOut.WriteObject(tree, tree->GetName());
+  }
 }
 
 void o2::utils::DebugStreamer::enableStream(const StreamFlags streamFlag)

Detectors/TPC/reconstruction/macro/createTPCSpaceChargeCorrection.C

@@ -54,6 +54,8 @@ std::unique_ptr<SC> spaceCharge;
 
 void getGlobalSpaceChargeCorrection(const int roc, double x, double y, double z,
                                     double& dx, double& dy, double& dz);
+void getLocalSpaceChargeCorrection(const int roc, int irow, double y, double z,
+                                   double& dx, double& dy, double& dz);
 void initSpaceCharge(const char* histoFileName, const char* histoName);
 
 void DumpFlatObjectToFile(const TPCFastTransform* obj, const char* file);
@@ -98,6 +100,56 @@ void createTPCSpaceChargeCorrection(
   }
 }
 
+/// Creates TPCFastTransform object for TPC space-charge correction, stores it in a file and provides a debug tree if requested
+/// \param outputFileName name of the output file to store the TPCFastTransform object in
+/// \param debug create debug tree comparing original corrections and spline interpolations from TPCFastTransform (1 = on the spline interpolation grid, 2 = on the original lookup table grid)
+void createTPCSpaceChargeCorrectionAnalytical(
+  const char* outputFileName = "tpctransform.root",
+  const int debug = 0)
+{
+  SC::setGrid(nZ, nR, nPhi);
+
+  // initialize space-charge object
+  spaceCharge = std::make_unique<SC>();
+
+  // default analytical formulas for distortions and corrections
+  AnalyticalDistCorr<double> formula;
+  formula.initDefault();
+  /*
+    set custom distortions and corrections here. e.g.
+    TFormula mDistZ{"mDlZ", "z * 11"};
+    formula.mDlZFormula = mDistZ;
+  */
+
+  spaceCharge->setDistortionsCorrectionsAnalytical(formula);
+  spaceCharge->setUseAnalyticalDistCorr(true);
+
+  // dumping the space-charge object to file to apply it during the TPC digitizer simulation
+  TFile fSC("distortions_analytical.root", "RECREATE");
+  spaceCharge->dumpAnalyticalCorrectionsDistortions(fSC);
+
+  TPCFastTransformHelperO2::instance()->setLocalSpaceChargeCorrection(getLocalSpaceChargeCorrection);
+
+  std::unique_ptr<TPCFastTransform> fastTransform(TPCFastTransformHelperO2::instance()->create(0));
+  fastTransform->writeToFile(outputFileName, "ccdb_object");
+
+  if (debug > 0) {
+    const o2::gpu::TPCFastTransformGeo& geo = fastTransform->getGeometry();
+    utils::TreeStreamRedirector pcstream(TString::Format("fastTransformUnitTest_debug%d_gridsize%d-%d-%d.root", debug, nPhi, nR, nZ).Data(), "recreate");
+    switch (debug) {
+      case 1:
+        debugInterpolation(pcstream, geo, fastTransform.get());
+        break;
+      case 2:
+        debugGridpoints(pcstream, geo, fastTransform.get());
+        break;
+      default:
+        printf("Debug option %d is not implemented. Debug tree will be empty.", debug);
+    }
+    pcstream.Close();
+  }
+}
+
 /// Initialize calculation of original correction lookup tables
 /// \param histoFileName path and name to the root file containing input space-charge density histograms
 /// \param histoName name of the input space-charge density histogram
@@ -131,6 +183,19 @@ void getGlobalSpaceChargeCorrection(const int roc, double x, double y, double z,
   spaceCharge->getCorrections(x, y, z, side, dx, dy, dz);
 }
 
+/// Function to get corrections from original lookup tables
+/// \param XYZ array with x, y and z position
+/// \param dXdYdZ array with correction dx, dy and dz
+void getLocalSpaceChargeCorrection(const int roc, int irow, double y, double z,
+                                   double& dx, double& dy, double& dz)
+{
+  Side side = roc < 18 ? Side::A : Side::C;
+  float x = o2::tpc::TPCFastTransformHelperO2::instance()->getGeometry().getRowInfo(irow).x;
+  dx = spaceCharge->getDistortionsCorrectionsAnalytical().getCorrectionsLX(x, y, z, side);
+  dy = spaceCharge->getDistortionsCorrectionsAnalytical().getCorrectionsLY(x, y, z, side);
+  dz = spaceCharge->getDistortionsCorrectionsAnalytical().getCorrectionsLZ(x, y, z, side);
+}
+
 /// Save TPCFastTransform to a file
 /// \param obj TPCFastTransform object to store
 /// \param file output file name

Detectors/TPC/simulation/src/Digitizer.cxx

@@ -163,6 +163,13 @@ void Digitizer::flush(std::vector<o2::tpc::Digit>& digits,
 {
   SAMPAProcessing& sampaProcessing = SAMPAProcessing::instance();
   mDigitContainer.fillOutputContainer(digits, labels, commonModeOutput, mSector, sampaProcessing.getTimeBinFromTime(mEventTime - mOutputDigitTimeOffset), mIsContinuous, finalFlush);
+  // flushing debug output to file
+  using Streamer = o2::utils::DebugStreamer;
+  if (Streamer::checkStream(o2::utils::StreamFlags::streamDistortionsSC)) {
+    if (((finalFlush && mIsContinuous) || (!mIsContinuous)) && mSpaceCharge) {
+      mSpaceCharge->flushStreamer();
+    }
+  }
 }
 
 void Digitizer::setUseSCDistortions(SC::SCDistortionType distortionType, const TH3* hisInitialSCDensity)
@@ -190,10 +197,13 @@ void Digitizer::setUseSCDistortions(TFile& finp)
   if (!mSpaceCharge) {
     mSpaceCharge = std::make_unique<SC>();
   }
-  mSpaceCharge->setGlobalDistortionsFromFile(finp, Side::A);
-  mSpaceCharge->setGlobalDistortionsFromFile(finp, Side::C);
-  mSpaceCharge->setGlobalCorrectionsFromFile(finp, Side::A);
-  mSpaceCharge->setGlobalCorrectionsFromFile(finp, Side::C);
+
+  // in case analytical distortions are loaded from file they are applied
+  mSpaceCharge->setAnalyticalCorrectionsDistortionsFromFile(finp);
+  if (!mSpaceCharge->getUseAnalyticalDistCorr()) {
+    mSpaceCharge->setGlobalDistortionsFromFile(finp, Side::A);
+    mSpaceCharge->setGlobalDistortionsFromFile(finp, Side::C);
+  }
 }
 
 void Digitizer::setStartTime(double time)