TPC: avoid having a direct dependency on RDataFrame

ktf · ktf · commit 9726160c1303 · 2026-06-24T21:30:45.000+02:00
diff --git a/Detectors/TPC/reconstruction/CMakeLists.txt b/Detectors/TPC/reconstruction/CMakeLists.txt
@@ -142,6 +142,7 @@ o2_add_test_root_macro(macro/createTPCSpaceChargeCorrection.C
                                              O2::CommonConstants
                                              O2::CommonUtils
                                              O2::TPCSpaceCharge
+                                             O2::TPCSpaceChargeIO
                        LABELS tpc)
 
 o2_add_test_root_macro(macro/findKrBoxCluster.C
diff --git a/Detectors/TPC/simulation/CMakeLists.txt b/Detectors/TPC/simulation/CMakeLists.txt
@@ -22,7 +22,7 @@ o2_add_library(TPCSimulation
                        src/SAMPAProcessing.cxx
                        src/IDCSim.cxx
                PUBLIC_LINK_LIBRARIES O2::DetectorsBase O2::SimulationDataFormat
-                                     O2::TPCBase O2::TPCSpaceCharge O2::TPCCalibration
+                                     O2::TPCBase O2::TPCSpaceCharge O2::TPCSpaceChargeIO O2::TPCCalibration
                                      ROOT::Physics)
 
 o2_target_root_dictionary(TPCSimulation
diff --git a/Detectors/TPC/spacecharge/CMakeLists.txt b/Detectors/TPC/spacecharge/CMakeLists.txt
@@ -19,9 +19,13 @@ o2_add_library(TPCSpaceCharge
                                      O2::Field
                                      Vc::Vc
                                      ROOT::Core
-                                     ROOT::ROOTDataFrame
                                      O2::DataFormatsParameters)
 
+o2_add_library(TPCSpaceChargeIO
+               SOURCES src/SpaceChargeIO.cxx
+               PUBLIC_LINK_LIBRARIES O2::TPCSpaceCharge
+                                     ROOT::ROOTDataFrame)
+
 
 o2_target_root_dictionary(TPCSpaceCharge
                           HEADERS include/TPCSpaceCharge/PoissonSolver.h
@@ -38,10 +42,12 @@ o2_target_root_dictionary(TPCSpaceCharge
 
 o2_add_test_root_macro(macro/calculateDistortionsCorrections.C
                        PUBLIC_LINK_LIBRARIES O2::TPCSpaceCharge
+                                             O2::TPCSpaceChargeIO
                        LABELS tpc COMPILE_ONLY)
 
 o2_add_test_root_macro(macro/createResidualDistortionObject.C
                       PUBLIC_LINK_LIBRARIES O2::TPCSpaceCharge
+                                            O2::TPCSpaceChargeIO
                                             O2::CommonUtils
                       LABELS tpc)
 
@@ -51,6 +57,7 @@ install(FILES macro/createSCHistosFromHits.C
 
 o2_add_test_root_macro(macro/createSCHistosFromHits.C
                       PUBLIC_LINK_LIBRARIES O2::TPCSpaceCharge
+                                            O2::TPCSpaceChargeIO
                                             O2::CommonUtils
                                             O2::TPCBase
                                             O2::TPCSimulation
diff --git a/Detectors/TPC/spacecharge/src/DataContainer3D.cxx b/Detectors/TPC/spacecharge/src/DataContainer3D.cxx
@@ -16,7 +16,6 @@
 #include "TPCBase/Mapper.h"
 #include "Framework/Logger.h"
 #include "TFile.h"
-#include "ROOT/RDataFrame.hxx"
 #include "TStopwatch.h"
 #include "TTree.h"
 
@@ -42,98 +41,7 @@ int DataContainer3D<DataT>::writeToFile(TFile& outf, const char* name) const
   return 0;
 }
 
-template <typename DataT>
-int DataContainer3D<DataT>::writeToFile(std::string_view file, std::string_view option, std::string_view name, const int nthreads) const
-{
-  // max number of floats per Entry
-  const size_t maxvalues = sizeof(float) * 1024 * 1024;
-
-  // total number of values to be stored
-  const size_t nsize = getNDataPoints();
-
-  // calculate number of entries in the tree and restrict if the number of values per threads exceeds max size
-  size_t entries = ((nsize / nthreads) > maxvalues) ? (nsize / maxvalues) : nthreads;
-
-  if (entries > nsize) {
-    entries = nsize;
-  }
-
-  // calculate numbers to store per entry
-  const size_t values_per_entry = nsize / entries;
 
-  // in case of remainder add additonal entry
-  const size_t values_lastEntry = nsize % entries;
-  if (values_lastEntry) {
-    entries += 1;
-  }
-
-  // in case EnableImplicitMT was already called with different number of threads, perform reset
-  if (ROOT::IsImplicitMTEnabled() && (ROOT::GetThreadPoolSize() != nthreads)) {
-    ROOT::DisableImplicitMT();
-  }
-  ROOT::EnableImplicitMT(nthreads);
-
-  // define dataframe which will be stored in the TTree
-  ROOT::RDataFrame dFrame(entries);
-
-  // define function which is used to fill the data frame
-  auto dfStore = dFrame.DefineSlotEntry(name, [&data = std::as_const(mData), entries, values_per_entry](unsigned int, ULong64_t entry) { return DataContainer3D<DataT>::getDataSlice(data, entries, values_per_entry, entry); });
-  dfStore = dfStore.Define("nz", [mZVertices = mZVertices]() { return mZVertices; });
-  dfStore = dfStore.Define("nr", [mRVertices = mRVertices]() { return mRVertices; });
-  dfStore = dfStore.Define("nphi", [mPhiVertices = mPhiVertices]() { return mPhiVertices; });
-
-  // define options of TFile
-  ROOT::RDF::RSnapshotOptions opt;
-  opt.fMode = option;
-  opt.fOverwriteIfExists = true; // overwrite if already exists
-
-  TStopwatch timer;
-  // note: first call has some overhead (~2s)
-  dfStore.Snapshot(name, file, {name.data(), "nz", "nr", "nphi"}, opt);
-  timer.Print("u");
-  return 0;
-}
-
-template <typename DataT>
-bool DataContainer3D<DataT>::initFromFile(std::string_view file, std::string_view name, const int nthreads)
-{
-  // in case EnableImplicitMT was already called with different number of threads, perform reset
-  if (ROOT::IsImplicitMTEnabled() && (ROOT::GetThreadPoolSize() != nthreads)) {
-    ROOT::DisableImplicitMT();
-  }
-  ROOT::EnableImplicitMT(nthreads);
-
-  // compare first the meta data (is the number of vertices the same)
-  // define data frame from imput file
-  ROOT::RDataFrame dFrame(name, file);
-
-  // compare vertices
-  auto comp = [mZVertices = mZVertices, mRVertices = mRVertices, mPhiVertices = mPhiVertices](const unsigned short nz, const unsigned short nr, const unsigned short nphi) {
-    if ((nz == mZVertices) && (nr == mRVertices) && (nphi == mPhiVertices)) {
-      return false;
-    }
-    return true;
-  };
-
-  auto count = dFrame.Filter(comp, {"nz", "nr", "nphi"}).Count();
-  if (*count != 0) {
-    LOGP(error, "Data from input file has different number of vertices! Found {} same vertices", *count);
-    return false;
-  }
-
-  // define lambda function which is used to copy the data
-  auto readData = [&mData = mData](const std::pair<long, std::vector<float>>& data) {
-    std::copy(data.second.begin(), data.second.end(), mData.begin() + data.first);
-  };
-
-  LOGP(info, "Reading {} from file {}", name, file);
-
-  // fill data from RDataFrame
-  TStopwatch timer;
-  dFrame.Foreach(readData, {name.data()});
-  timer.Print("u");
-  return true;
-}
 
 /// set values from file
 template <typename DataT>
@@ -215,18 +123,6 @@ void DataContainer3D<DataT>::print() const
   LOGP(info, "{} \n \n", stream.str());
 }
 
-template <typename DataT>
-auto DataContainer3D<DataT>::getDataSlice(const std::vector<DataT>& data, size_t entries, const size_t values_per_entry, ULong64_t entry)
-{
-  const long indStart = entry * values_per_entry;
-  if (entry < (entries - 1)) {
-    return std::pair(indStart, std::vector<float>(data.begin() + indStart, data.begin() + indStart + values_per_entry));
-  } else if (entry == (entries - 1)) {
-    // last entry might have different number of values. just copy the rest...
-    return std::pair(indStart, std::vector<float>(data.begin() + indStart, data.end()));
-  }
-  return std::pair(indStart, std::vector<float>());
-};
 
 template <typename DataT>
 DataContainer3D<DataT>& DataContainer3D<DataT>::operator*=(const DataT value)
@@ -321,82 +217,6 @@ void DataContainer3D<DataT>::setGrid(unsigned short nZ, unsigned short nR, unsig
   }
 }
 
-template <typename DataT>
-void DataContainer3D<DataT>::dumpSlice(std::string_view treename, std::string_view fileIn, std::string_view fileOut, std::string_view option, std::pair<unsigned short, unsigned short> rangeiR, std::pair<unsigned short, unsigned short> rangeiZ, std::pair<unsigned short, unsigned short> rangeiPhi, const int nthreads)
-{
-  if (ROOT::IsImplicitMTEnabled() && (ROOT::GetThreadPoolSize() != nthreads)) {
-    ROOT::DisableImplicitMT();
-  }
-  ROOT::EnableImplicitMT(nthreads);
-  ROOT::RDataFrame dFrame(treename, fileIn);
-
-  auto df = dFrame.Define("slice", [rangeiZ, rangeiR, rangeiPhi](const std::pair<long, std::vector<float>>& values, unsigned short nz, unsigned short nr, unsigned short nphi) {
-    std::vector<size_t> ir;
-    std::vector<size_t> iphi;
-    std::vector<size_t> iz;
-    std::vector<float> r;
-    std::vector<float> phi;
-    std::vector<float> z;
-    std::vector<float> vals;
-    std::vector<size_t> globalIdx;
-    std::vector<LocalPosition3D> lPos;
-    const auto nvalues = values.second.size();
-    ir.reserve(nvalues);
-    iphi.reserve(nvalues);
-    iz.reserve(nvalues);
-    r.reserve(nvalues);
-    phi.reserve(nvalues);
-    z.reserve(nvalues);
-    vals.reserve(nvalues);
-    lPos.reserve(nvalues);
-    globalIdx.reserve(nvalues);
-    for (size_t i = 0; i < nvalues; ++i) {
-      const size_t idx = values.first + i;
-      const auto iZTmp = o2::tpc::DataContainer3D<float>::getIndexZ(idx, nz, nr, nphi);
-      if ((rangeiZ.first < rangeiZ.second) && ((iZTmp < rangeiZ.first) || (iZTmp > rangeiZ.second))) {
-        continue;
-      }
-
-      const auto iRTmp = o2::tpc::DataContainer3D<float>::getIndexR(idx, nz, nr, nphi);
-      if ((rangeiR.first < rangeiR.second) && ((iRTmp < rangeiR.first) || (iRTmp > rangeiR.second))) {
-        continue;
-      }
-
-      const auto iPhiTmp = o2::tpc::DataContainer3D<float>::getIndexPhi(idx, nz, nr, nphi);
-      if ((rangeiPhi.first < rangeiPhi.second) && ((iPhiTmp < rangeiPhi.first) || (iPhiTmp > rangeiPhi.second))) {
-        continue;
-      }
-
-      const float rTmp = o2::tpc::GridProperties<float>::getRMin() + o2::tpc::GridProperties<float>::getGridSpacingR(nr) * iRTmp;
-      const float zTmp = o2::tpc::GridProperties<float>::getZMin() + o2::tpc::GridProperties<float>::getGridSpacingZ(nz) * iZTmp;
-      const float phiTmp = o2::tpc::GridProperties<float>::getPhiMin() + o2::tpc::GridProperties<float>::getGridSpacingPhi(nphi) * (MGParameters::normalizeGridToNSector / double(SECTORSPERSIDE)) * iPhiTmp;
-
-      const float x = rTmp * std::cos(phiTmp);
-      const float y = rTmp * std::sin(phiTmp);
-      const LocalPosition3D pos(x, y, zTmp);
-      unsigned char secNum = std::floor(phiTmp / SECPHIWIDTH);
-      Sector sector(secNum + (pos.Z() < 0) * SECTORSPERSIDE);
-      LocalPosition3D lPosTmp = Mapper::GlobalToLocal(pos, sector);
-
-      lPos.emplace_back(lPosTmp);
-      ir.emplace_back(iRTmp);
-      iphi.emplace_back(iPhiTmp);
-      iz.emplace_back(iZTmp);
-      r.emplace_back(rTmp);
-      phi.emplace_back(phiTmp);
-      z.emplace_back(zTmp);
-      vals.emplace_back(values.second[i]);
-      globalIdx.emplace_back(idx);
-    }
-    return std::make_tuple(vals, iz, ir, iphi, z, r, phi, lPos, globalIdx);
-  },
-                          {treename.data(), "nz", "nr", "nphi"});
-
-  // define options of TFile
-  ROOT::RDF::RSnapshotOptions opt;
-  opt.fMode = option;
-  df.Snapshot(treename, fileOut, {"slice"}, opt);
-}
 
 template <typename DataT>
 DataT DataContainer3D<DataT>::interpolate(const DataT z, const DataT r, const DataT phi, const o2::tpc::RegularGrid3D<DataT>& grid) const
@@ -405,93 +225,6 @@ DataT DataContainer3D<DataT>::interpolate(const DataT z, const DataT r, const Da
   return interpolator(z, r, phi);
 }
 
-template <typename DataT>
-void DataContainer3D<DataT>::dumpInterpolation(std::string_view treename, std::string_view fileIn, std::string_view fileOut, std::string_view option, std::pair<float, float> rangeR, std::pair<float, float> rangeZ, std::pair<float, float> rangePhi, const int nR, const int nZ, const int nPhi, const int nthreads)
-{
-  if (ROOT::IsImplicitMTEnabled() && (ROOT::GetThreadPoolSize() != nthreads)) {
-    ROOT::DisableImplicitMT();
-  }
-  ROOT::EnableImplicitMT(nthreads);
-  ROOT::RDataFrame dFrame(nPhi);
-
-  // get vertices for input TTree which is needed to define the grid for interpolation
-  unsigned short nr, nz, nphi;
-  if (!getVertices(treename, fileIn, nr, nz, nphi)) {
-    return;
-  }
-
-  // load data from input TTree
-  DataContainer3D<DataT> data;
-  data.setGrid(nz, nr, nphi, true);
-  data.initFromFile(fileIn, treename, nthreads);
-
-  // define grid for interpolation
-  using GridProp = GridProperties<DataT>;
-  const RegularGrid3D<DataT> mGrid3D(GridProp::ZMIN, GridProp::RMIN, GridProp::PHIMIN, GridProp::getGridSpacingZ(nz), GridProp::getGridSpacingR(nr), o2::tpc::GridProperties<float>::getGridSpacingPhi(nphi) * (MGParameters::normalizeGridToNSector / double(SECTORSPERSIDE)), ParamSpaceCharge{nr, nz, nphi});
-
-  auto interpolate = [&mGrid3D = std::as_const(mGrid3D), &data = std::as_const(data), rangeR, rangeZ, rangePhi, nR, nZ, nPhi](unsigned int, ULong64_t iPhi) {
-    std::vector<size_t> ir;
-    std::vector<size_t> iphi;
-    std::vector<size_t> iz;
-    std::vector<float> r;
-    std::vector<float> phi;
-    std::vector<float> z;
-    std::vector<float> vals;
-    std::vector<size_t> globalIdx;
-    std::vector<LocalPosition3D> lPos;
-    const auto nvalues = nR * nZ;
-    ir.reserve(nvalues);
-    iphi.reserve(nvalues);
-    iz.reserve(nvalues);
-    r.reserve(nvalues);
-    phi.reserve(nvalues);
-    z.reserve(nvalues);
-    vals.reserve(nvalues);
-    lPos.reserve(nvalues);
-    globalIdx.reserve(nvalues);
-
-    const float rSpacing = (rangeR.second - rangeR.first) / (nR - 1);
-    const float zSpacing = (rangeZ.second - rangeZ.first) / (nZ - 1);
-    const float phiSpacing = (rangePhi.second - rangePhi.first) / (nPhi - 1);
-    const DataT phiPos = rangePhi.first + iPhi * phiSpacing;
-    // loop over grid and interpolate values
-    for (int iR = 0; iR < nR; ++iR) {
-      const DataT rPos = rangeR.first + iR * rSpacing;
-      for (int iZ = 0; iZ < nZ; ++iZ) {
-        const size_t idx = (iZ + nZ * (iR + iPhi * nR)); // unique index to Build index with other friend TTrees
-        const DataT zPos = rangeZ.first + iZ * zSpacing;
-        ir.emplace_back(iR);
-        iphi.emplace_back(iPhi);
-        iz.emplace_back(iZ);
-        r.emplace_back(rPos);
-        phi.emplace_back(phiPos);
-        z.emplace_back(zPos);
-        vals.emplace_back(data.interpolate(zPos, rPos, phiPos, mGrid3D)); // interpolated values
-        globalIdx.emplace_back(idx);
-        const float x = rPos * std::cos(phiPos);
-        const float y = rPos * std::sin(phiPos);
-        const LocalPosition3D pos(x, y, zPos);
-        unsigned char secNum = std::floor(phiPos / SECPHIWIDTH); // TODO CHECK THIS
-        Sector sector(secNum + (pos.Z() < 0) * SECTORSPERSIDE);
-        LocalPosition3D lPosTmp = Mapper::GlobalToLocal(pos, sector);
-        lPos.emplace_back(lPosTmp);
-      }
-    }
-    return std::make_tuple(vals, iz, ir, iphi, z, r, phi, lPos, globalIdx);
-  };
-
-  // define RDataFrame entry
-  auto dfStore = dFrame.DefineSlotEntry(treename, interpolate);
-
-  // define options of TFile
-  ROOT::RDF::RSnapshotOptions opt;
-  opt.fMode = option;
-
-  TStopwatch timer;
-  // note: first call has some overhead (~2s)
-  dfStore.Snapshot(treename, fileOut, {treename.data()}, opt);
-  timer.Print("u");
-}
 
 template <typename DataT>
 bool DataContainer3D<DataT>::getVertices(std::string_view treename, std::string_view fileIn, unsigned short& nR, unsigned short& nZ, unsigned short& nPhi)
diff --git a/Detectors/TPC/spacecharge/src/SpaceCharge.cxx b/Detectors/TPC/spacecharge/src/SpaceCharge.cxx
diff --git a/Detectors/TPC/spacecharge/src/SpaceChargeIO.cxx b/Detectors/TPC/spacecharge/src/SpaceChargeIO.cxx
