Skip to content

TPC: avoid having a direct dependency on RDataFrame #15551

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ktf wants to merge 1 commit into
base: dev
Choose a base branch
from
+1,093 −960
Open

TPC: avoid having a direct dependency on RDataFrame #15551

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions Detectors/TPC/reconstruction/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion Detectors/TPC/simulation/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
9 changes: 8 additions & 1 deletion Detectors/TPC/spacecharge/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)

Expand All @@ -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
Expand Down
267 changes: 0 additions & 267 deletions Detectors/TPC/spacecharge/src/DataContainer3D.cxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand All @@ -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>
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -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)
Expand Down
Loading
Loading