19#include "ROOT/RDataFrame.hxx"
20#include "TStopwatch.h"
29template <
typename DataT>
30template <
typename DataTOut>
33 if (outf.IsZombie()) {
34 LOGP(error,
"Failed to write to file: {}", outf.GetName());
39 containerTmp.
getData() = std::vector<DataTOut>(mData.begin(), mData.end());
45template <
typename DataT>
49 const size_t maxvalues =
sizeof(float) * 1024 * 1024;
52 const size_t nsize = getNDataPoints();
55 size_t entries = ((nsize / nthreads) > maxvalues) ? (nsize / maxvalues) : nthreads;
57 if (entries > nsize) {
62 const size_t values_per_entry = nsize / entries;
65 const size_t values_lastEntry = nsize % entries;
66 if (values_lastEntry) {
71 if (ROOT::IsImplicitMTEnabled() && (ROOT::GetThreadPoolSize() != nthreads)) {
72 ROOT::DisableImplicitMT();
74 ROOT::EnableImplicitMT(nthreads);
77 ROOT::RDataFrame dFrame(entries);
81 dfStore = dfStore.Define(
"nz", [mZVertices = mZVertices]() {
return mZVertices; });
82 dfStore = dfStore.Define(
"nr", [mRVertices = mRVertices]() {
return mRVertices; });
83 dfStore = dfStore.Define(
"nphi", [mPhiVertices = mPhiVertices]() {
return mPhiVertices; });
86 ROOT::RDF::RSnapshotOptions opt;
88 opt.fOverwriteIfExists =
true;
92 dfStore.Snapshot(
name,
file, {
name.data(),
"nz",
"nr",
"nphi"}, opt);
97template <
typename DataT>
101 if (ROOT::IsImplicitMTEnabled() && (ROOT::GetThreadPoolSize() != nthreads)) {
102 ROOT::DisableImplicitMT();
104 ROOT::EnableImplicitMT(nthreads);
108 ROOT::RDataFrame dFrame(
name,
file);
111 auto comp = [mZVertices = mZVertices, mRVertices = mRVertices, mPhiVertices = mPhiVertices](
const unsigned short nz,
const unsigned short nr,
const unsigned short nphi) {
112 if ((nz == mZVertices) && (nr == mRVertices) && (nphi == mPhiVertices)) {
118 auto count = dFrame.Filter(comp, {
"nz",
"nr",
"nphi"}).Count();
120 LOGP(error,
"Data from input file has different number of vertices! Found {} same vertices", *
count);
125 auto readData = [&mData = mData](
const std::pair<long, std::vector<float>>&
data) {
126 std::copy(
data.second.begin(),
data.second.end(), mData.begin() +
data.first);
129 LOGP(info,
"Reading {} from file {}",
name,
file);
139template <
typename DataT>
140template <
typename DataTIn>
143 if (inpf.IsZombie()) {
144 LOGP(error,
"Failed to read from file: {}", inpf.GetName());
151 LOGP(error,
"Failed to load {} from {}",
name, inpf.GetName());
155 if (mZVertices != dataCont->getNZ() || mRVertices != dataCont->getNR() || mPhiVertices != dataCont->getNPhi()) {
156 LOGP(error,
"Data from input file has different definition of vertices!");
157 LOGP(error,
"set vertices before creating the sc object to: SpaceCharge<>::setGrid({}, {}, {})", dataCont->getNZ(), dataCont->getNR(), dataCont->getNPhi());
162 mData = std::vector<DataT>(dataCont->getData().begin(), dataCont->getData().end());
167template <
typename DataT>
170 if (inpf.IsZombie()) {
171 LOGP(error,
"Failed to read from file {}", inpf.GetName());
178 LOGP(error,
"Failed to load {} from {}",
name, inpf.GetName());
184template <
typename DataT>
189 auto&&
w = std::setw(9);
192 for (
unsigned int iz = 0; iz < mPhiVertices; ++iz) {
195 stream <<
"⎡" <<
w << (*this)(0, 0, iz);
196 for (
unsigned int ix = 1; ix < mZVertices; ++ix) {
197 stream <<
", " <<
w << (*this)(ix, 0, iz);
201 for (
unsigned int iy = 1; iy < mRVertices - 1; ++iy) {
202 stream <<
"⎢" <<
w << (*this)(0, iy, iz);
203 for (
unsigned int ix = 1; ix < mZVertices; ++ix) {
204 stream <<
", " <<
w << (*this)(ix, iy, iz);
209 stream <<
"⎣" <<
w << (*this)(0, mRVertices - 1, iz);
210 for (
unsigned int ix = 1; ix < mZVertices; ++ix) {
211 stream <<
", " <<
w << (*this)(ix, mRVertices - 1, iz);
215 LOGP(info,
"{} \n \n",
stream.str());
218template <
typename DataT>
221 const long indStart =
entry * values_per_entry;
222 if (
entry < (entries - 1)) {
223 return std::pair(indStart, std::vector<float>(
data.begin() + indStart,
data.begin() + indStart + values_per_entry));
224 }
else if (
entry == (entries - 1)) {
226 return std::pair(indStart, std::vector<float>(
data.begin() + indStart,
data.end()));
228 return std::pair(indStart, std::vector<float>());
231template <
typename DataT>
234 std::transform(mData.begin(), mData.end(), mData.begin(), [
value =
value](
auto&
val) { return val * value; });
238template <
typename DataT>
241 std::transform(mData.begin(), mData.end(),
other.mData.begin(), mData.begin(), std::plus<>());
245template <
typename DataT>
248 std::transform(mData.begin(), mData.end(),
other.mData.begin(), mData.begin(), std::minus<>());
252template <
typename DataT>
255 std::transform(mData.begin(), mData.end(),
other.mData.begin(), mData.begin(), std::multiplies<>());
259template <
typename DataT>
262 const size_t iphi =
index / (nz * nr);
263 index -= (iphi * nz * nr);
264 const size_t iz =
index % nz;
268template <
typename DataT>
271 const size_t iphi =
index / (nz * nr);
272 index -= (iphi * nz * nr);
277template <
typename DataT>
280 return index / (nz * nr);
283template <
typename DataT>
286 tree->SetAlias(
"ir",
"o2::tpc::DataContainer3D<float>::getIndexR(first + Iteration$, nz, nr, nphi)");
287 tree->SetAlias(
"iz",
"o2::tpc::DataContainer3D<float>::getIndexZ(first + Iteration$, nz, nr, nphi)");
288 tree->SetAlias(
"iphi",
"o2::tpc::DataContainer3D<float>::getIndexPhi(first + Iteration$, nz, nr, nphi)");
289 tree->SetAlias(
"r",
"o2::tpc::GridProperties<float>::getRMin() + o2::tpc::GridProperties<float>::getGridSpacingR(nr) * ir");
290 tree->SetAlias(
"z",
"o2::tpc::GridProperties<float>::getZMin() + o2::tpc::GridProperties<float>::getGridSpacingZ(nz) * iz");
291 tree->SetAlias(
"phi",
"o2::tpc::GridProperties<float>::getPhiMin() + o2::tpc::GridProperties<float>::getGridSpacingPhi(nphi) * iphi");
294template <
typename DataT>
298 tree->SetAlias(
"val",
"_0");
301 tree->SetAlias(
"iz",
"_1");
302 tree->SetAlias(
"ir",
"_2");
303 tree->SetAlias(
"iphi",
"_3");
304 tree->SetAlias(
"z",
"_4");
305 tree->SetAlias(
"r",
"_5");
306 tree->SetAlias(
"phi",
"_6");
307 tree->SetAlias(
"lpos",
"_7");
308 tree->SetAlias(
"lx",
"lpos.fCoordinates.fX");
309 tree->SetAlias(
"ly",
"lpos.fCoordinates.fY");
310 tree->SetAlias(
"index",
"_8");
313template <
typename DataT>
320 mData.resize(nZ * nR *
static_cast<size_t>(nPhi));
324template <
typename DataT>
325void 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)
327 if (ROOT::IsImplicitMTEnabled() && (ROOT::GetThreadPoolSize() != nthreads)) {
328 ROOT::DisableImplicitMT();
330 ROOT::EnableImplicitMT(nthreads);
331 ROOT::RDataFrame dFrame(treename, fileIn);
333 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) {
335 std::vector<size_t>
ir;
336 std::vector<size_t> iphi;
337 std::vector<size_t> iz;
338 std::vector<float>
r;
339 std::vector<float> phi;
340 std::vector<float>
z;
341 std::vector<float> vals;
342 std::vector<size_t> globalIdx;
343 std::vector<LocalPosition3D> lPos;
344 const auto nvalues =
values.second.size();
346 iphi.reserve(nvalues);
349 phi.reserve(nvalues);
351 vals.reserve(nvalues);
352 lPos.reserve(nvalues);
353 globalIdx.reserve(nvalues);
354 for (
size_t i = 0;
i < nvalues; ++
i) {
355 const size_t idx =
values.first +
i;
357 if ((rangeiZ.first < rangeiZ.second) && ((iZTmp < rangeiZ.first) || (iZTmp > rangeiZ.second))) {
362 if ((rangeiR.first < rangeiR.second) && ((iRTmp < rangeiR.first) || (iRTmp > rangeiR.second))) {
367 if ((rangeiPhi.first < rangeiPhi.second) && ((iPhiTmp < rangeiPhi.first) || (iPhiTmp > rangeiPhi.second))) {
375 const float x = rTmp * std::cos(phiTmp);
376 const float y = rTmp * std::sin(phiTmp);
378 unsigned char secNum = simOneSectorOnly ? 0 : std::floor(phiTmp /
SECPHIWIDTH);
382 lPos.emplace_back(lPosTmp);
383 ir.emplace_back(iRTmp);
384 iphi.emplace_back(iPhiTmp);
385 iz.emplace_back(iZTmp);
386 r.emplace_back(rTmp);
387 phi.emplace_back(phiTmp);
388 z.emplace_back(zTmp);
389 vals.emplace_back(
values.second[
i]);
390 globalIdx.emplace_back(idx);
392 return std::make_tuple(vals, iz,
ir, iphi,
z,
r, phi, lPos, globalIdx);
394 {treename.data(),
"nz",
"nr",
"nphi"});
397 ROOT::RDF::RSnapshotOptions opt;
399 df.Snapshot(treename, fileOut, {
"slice"}, opt);
402template <
typename DataT>
406 return interpolator(
z,
r, phi);
409template <
typename DataT>
410void 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)
412 if (ROOT::IsImplicitMTEnabled() && (ROOT::GetThreadPoolSize() != nthreads)) {
413 ROOT::DisableImplicitMT();
415 ROOT::EnableImplicitMT(nthreads);
416 ROOT::RDataFrame dFrame(nPhi);
419 unsigned short nr, nz, nphi;
420 if (!getVertices(treename, fileIn, nr, nz, nphi)) {
427 data.initFromFile(fileIn, treename, nthreads);
433 auto interpolate = [&mGrid3D = std::as_const(mGrid3D), &
data = std::as_const(
data), rangeR, rangeZ, rangePhi, nR, nZ, nPhi](
unsigned int, ULong64_t iPhi) {
435 std::vector<size_t>
ir;
436 std::vector<size_t> iphi;
437 std::vector<size_t> iz;
438 std::vector<float>
r;
439 std::vector<float> phi;
440 std::vector<float>
z;
441 std::vector<float> vals;
442 std::vector<size_t> globalIdx;
443 std::vector<LocalPosition3D> lPos;
444 const auto nvalues = nR * nZ;
446 iphi.reserve(nvalues);
449 phi.reserve(nvalues);
451 vals.reserve(nvalues);
452 lPos.reserve(nvalues);
453 globalIdx.reserve(nvalues);
455 const float rSpacing = (rangeR.second - rangeR.first) / (nR - 1);
456 const float zSpacing = (rangeZ.second - rangeZ.first) / (nZ - 1);
457 const float phiSpacing = (rangePhi.second - rangePhi.first) / (nPhi - 1);
458 const DataT phiPos = rangePhi.first + iPhi * phiSpacing;
460 for (
int iR = 0; iR < nR; ++iR) {
461 const DataT rPos = rangeR.first + iR * rSpacing;
462 for (
int iZ = 0; iZ < nZ; ++iZ) {
463 const size_t idx = (iZ + nZ * (iR + iPhi * nR));
464 const DataT zPos = rangeZ.first + iZ * zSpacing;
466 iphi.emplace_back(iPhi);
468 r.emplace_back(rPos);
469 phi.emplace_back(phiPos);
470 z.emplace_back(zPos);
471 vals.emplace_back(
data.interpolate(zPos, rPos, phiPos, mGrid3D));
472 globalIdx.emplace_back(idx);
473 const float x = rPos * std::cos(phiPos);
474 const float y = rPos * std::sin(phiPos);
476 unsigned char secNum = simOneSectorOnly ? 0 : std::floor(phiPos /
SECPHIWIDTH);
479 lPos.emplace_back(lPosTmp);
482 return std::make_tuple(vals, iz,
ir, iphi,
z,
r, phi, lPos, globalIdx);
486 auto dfStore = dFrame.DefineSlotEntry(treename, interpolate);
489 ROOT::RDF::RSnapshotOptions opt;
494 dfStore.Snapshot(treename, fileOut, {treename.data()}, opt);
498template <
typename DataT>
501 TFile fTmp(fileIn.data(),
"READ");
502 TTree*
tree = (TTree*)fTmp.Get(treename.data());
504 LOGP(warning,
"Tree {} not found in input file {}", treename, fileIn);
507 tree->SetBranchAddress(
"nz", &nZ);
508 tree->SetBranchAddress(
"nr", &nR);
509 tree->SetBranchAddress(
"nphi", &nPhi);
This class provides a simple method to store values on a large 3-Dim grid with ROOT io functionality.
Definition of RegularGrid3D class.
Definition of TriCubic class.
static LocalPosition3D GlobalToLocal(const GlobalPosition3D &pos, const double alpha)
GLuint const GLchar * name
GLsizei const GLfloat * value
GLenum GLsizei GLsizei GLint * values
GLubyte GLubyte GLubyte GLubyte w
GLdouble GLdouble GLdouble z
Global TPC definitions and constants.
constexpr double SECPHIWIDTH
constexpr unsigned char SECTORSPERSIDE
void readData(o2::tpc::GBTFrameContainer &container, std::vector< std::ofstream * > &outfiles, int &run, int &done)
const auto & getData() const
static size_t getIndexR(size_t index, const int nz, const int nr, const int nphi)
static void 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=1)
void print() const
print the matrix
static void 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=1)
bool initFromFile(TFile &inpf, const char *name="data")
set values from file
int writeToFile(TFile &outf, const char *name="data") const
void setGrid(unsigned short nZ, unsigned short nR, unsigned short nPhi, const bool resize)
set the grid points
static void setAliasesForDump(TTree *tree)
static void setAliases(TTree *tree)
DataT interpolate(const DataT z, const DataT r, const DataT phi, const o2::tpc::RegularGrid3D< DataT > &grid) const
DataContainer3D< DataT > & operator*=(const DataT value)
operator overload
static size_t getIndexZ(size_t index, const int nz, const int nr, const int nphi)
static size_t getIndexPhi(size_t index, const int nz, const int nr, const int nphi)
static bool getVertices(std::string_view treename, std::string_view fileIn, unsigned short &nR, unsigned short &nZ, unsigned short &nPhi)
static DataContainer3D< DataT > * loadFromFile(TFile &inpf, const char *name="data")
get pointer to object from file (deprecated!)
DataContainer3D< DataT > & operator-=(const DataContainer3D< DataT > &other)
DataContainer3D< DataT > & operator+=(const DataContainer3D< DataT > &other)
static constexpr DataT getRMin()
static constexpr DataT getZMin()
static constexpr DataT getGridSpacingR(const unsigned int nR)
static constexpr DataT getPhiMin()
static constexpr DataT getGridSpacingZ(const unsigned int nZ)
static constexpr DataT getGridSpacingPhi(const unsigned int nPhi)
static bool normalizeGridToOneSector
the grid in phi direction is squashed from 2 Pi to (2 Pi / SECTORSPERSIDE). This can used to get the ...
VectorOfTObjectPtrs other
o2::InteractionRecord ir(0, 0)
std::unique_ptr< TTree > tree((TTree *) flIn.Get(std::string(o2::base::NameConf::CTFTREENAME).c_str()))