diff --git a/tree/dataframe/inc/ROOT/RDF/ActionHelpers.hxx b/tree/dataframe/inc/ROOT/RDF/ActionHelpers.hxx index c64870a0e0c48..f6fe6121d7cd3 100644 --- a/tree/dataframe/inc/ROOT/RDF/ActionHelpers.hxx +++ b/tree/dataframe/inc/ROOT/RDF/ActionHelpers.hxx @@ -513,6 +513,15 @@ public: void Initialize() {} void InitTask(TTreeReader *, unsigned int) {} + template + void ExecLoop(unsigned int slot, std::size_t elements, Iterators... its) + { + for (std::size_t i = 0; i < elements; i++) { + Exec(slot, *its...); + (std::advance(its, 1), ...); + } + } + template void ExecWithWeight(unsigned int slot, const std::tuple &columnValues, std::index_sequence) @@ -526,7 +535,22 @@ public: template void Exec(unsigned int slot, const ColumnTypes &...columnValues) { - if constexpr (WithWeight) { + if constexpr (std::disjunction_v...>) { + constexpr std::array isContainer{IsDataContainer::value...}; + constexpr std::size_t firstContainerIdx = FindIdxTrue(isContainer); + std::array sizes = {{GetSize(columnValues)...}}; + std::size_t elements = 0; + for (std::size_t i = 0; i < isContainer.size(); i++) { + if (isContainer[i]) { + if (i == firstContainerIdx) { + elements = sizes[i]; + } else if (elements != sizes[i]) { + throw std::runtime_error("Cannot fill values in containers of different sizes."); + } + } + } + ExecLoop(slot, elements, MakeBegin(columnValues)...); + } else if constexpr (WithWeight) { auto t = std::forward_as_tuple(columnValues...); ExecWithWeight(slot, t, std::make_index_sequence()); } else { @@ -566,6 +590,15 @@ public: void Initialize() {} void InitTask(TTreeReader *, unsigned int) {} + template + void ExecLoop(unsigned int slot, std::size_t elements, Iterators... its) + { + for (std::size_t i = 0; i < elements; i++) { + Exec(slot, *its...); + (std::advance(its, 1), ...); + } + } + template void ExecWithWeight(const std::tuple &columnValues, std::index_sequence) { @@ -576,9 +609,24 @@ public: } template - void Exec(unsigned int, const ColumnTypes &...columnValues) + void Exec(unsigned int slot, const ColumnTypes &...columnValues) { - if constexpr (WithWeight) { + if constexpr (std::disjunction_v...>) { + constexpr std::array isContainer{IsDataContainer::value...}; + constexpr std::size_t firstContainerIdx = FindIdxTrue(isContainer); + std::array sizes = {{GetSize(columnValues)...}}; + std::size_t elements = 0; + for (std::size_t i = 0; i < isContainer.size(); i++) { + if (isContainer[i]) { + if (i == firstContainerIdx) { + elements = sizes[i]; + } else if (elements != sizes[i]) { + throw std::runtime_error("Cannot fill values in containers of different sizes."); + } + } + } + ExecLoop(slot, elements, MakeBegin(columnValues)...); + } else if constexpr (WithWeight) { auto t = std::forward_as_tuple(columnValues...); ExecWithWeight(t, std::make_index_sequence()); } else { diff --git a/tree/dataframe/test/dataframe_hist.cxx b/tree/dataframe/test/dataframe_hist.cxx index f5fe396b50de3..8d53cf44c612d 100644 --- a/tree/dataframe/test/dataframe_hist.cxx +++ b/tree/dataframe/test/dataframe_hist.cxx @@ -8,6 +8,7 @@ #include #include #include +#include #include #include @@ -455,6 +456,280 @@ TEST_P(RDFHist, WeightInvalidNumberOfArgumentsJit) EXPECT_THROW(dfXW.Hist(engine, {"x", "x"}, "w"), std::invalid_argument); } +TEST_P(RDFHist, Container) +{ + RDataFrame df(10); + auto dfX = df.Define("x", [](ULong64_t e) { return ROOT::RVecD{e + 5.5}; }, {"rdfentry_"}); + + const RRegularAxis axis(10, {5.0, 15.0}); + auto hist = dfX.Hist({axis}, {"x"}); + EXPECT_EQ(hist->GetNEntries(), 10); + for (auto index : axis.GetNormalRange()) { + EXPECT_EQ(hist->GetBinContent(index), 1.0); + } +} + +TEST_P(RDFHist, ContainerEmpty) +{ + RDataFrame df(10); + auto dfX = df.Define("x", [] { return ROOT::RVecD{}; }); + + const RRegularAxis axis(10, {5.0, 15.0}); + auto hist = dfX.Hist({axis}, {"x"}); + EXPECT_EQ(hist->GetNEntries(), 0); +} + +TEST_P(RDFHist, ContainerElements) +{ + RDataFrame df(10); + auto dfX = df.Define("x", [](ULong64_t e) { return ROOT::RVecD{e + 5.5, e + 15.5}; }, {"rdfentry_"}); + + const RRegularAxis axis(20, {5.0, 25.0}); + auto hist = dfX.Hist({axis}, {"x"}); + EXPECT_EQ(hist->GetNEntries(), 20); + for (auto index : axis.GetNormalRange()) { + EXPECT_EQ(hist->GetBinContent(index), 1.0); + } +} + +TEST_P(RDFHist, ContainerNested) +{ + RDataFrame df(10); + auto dfX = df.Define("x", [](ULong64_t e) { return ROOT::RVec{{e + 5.5}, {e + 15.5}}; }, {"rdfentry_"}); + + const RRegularAxis axis(20, {5.0, 25.0}); + auto hist = dfX.Hist>({axis}, {"x"}); + EXPECT_EQ(hist->GetNEntries(), 20); + for (auto index : axis.GetNormalRange()) { + EXPECT_EQ(hist->GetBinContent(index), 1.0); + } +} + +TEST_P(RDFHist, ContainerJit) +{ + RDataFrame df(10); + auto dfX = df.Define("x", "ROOT::RVecD{rdfentry_ + 5.5}"); + + const RRegularAxis axis(10, {5.0, 15.0}); + auto hist = dfX.Hist({axis}, {"x"}); + EXPECT_EQ(hist->GetNEntries(), 10); + for (auto index : axis.GetNormalRange()) { + EXPECT_EQ(hist->GetBinContent(index), 1.0); + } +} + +TEST_P(RDFHist, ContainerMultiDim) +{ + RDataFrame df(10); + auto dfXY = df.Define("x", [](ULong64_t e) { return ROOT::RVecD{e + 5.5}; }, {"rdfentry_"}) + .Define("y", [](ULong64_t e) { return ROOT::RVecD{2 * e + 0.5}; }, {"rdfentry_"}); + + const RRegularAxis regularAxis(10, {5.0, 15.0}); + static constexpr std::size_t BinsY = 20; + std::vector bins; + for (std::size_t i = 0; i < BinsY; i++) { + bins.push_back(i); + } + bins.push_back(BinsY); + const RVariableBinAxis variableBinAxis(bins); + + auto hist = + dfXY.Hist({regularAxis, variableBinAxis}, {"x", "y"}); + EXPECT_EQ(hist->GetNEntries(), 10); + for (auto x : regularAxis.GetNormalRange()) { + for (auto y : variableBinAxis.GetNormalRange()) { + if (2 * x.GetIndex() == y.GetIndex()) { + EXPECT_EQ(hist->GetBinContent(x, y), 1.0); + } else { + EXPECT_EQ(hist->GetBinContent(x, y), 0.0); + } + } + } +} + +TEST_P(RDFHist, ContainerMultiDimBroadcast) +{ + RDataFrame df(10); + auto dfXY = df.Define("x", [](ULong64_t e) { return e + 5.5; }, {"rdfentry_"}) + .Define("y", [](ULong64_t e) { return ROOT::RVecD{2 * e + 0.5}; }, {"rdfentry_"}); + + const RRegularAxis regularAxis(10, {5.0, 15.0}); + static constexpr std::size_t BinsY = 20; + std::vector bins; + for (std::size_t i = 0; i < BinsY; i++) { + bins.push_back(i); + } + bins.push_back(BinsY); + const RVariableBinAxis variableBinAxis(bins); + + auto hist = dfXY.Hist({regularAxis, variableBinAxis}, {"x", "y"}); + EXPECT_EQ(hist->GetNEntries(), 10); + for (auto x : regularAxis.GetNormalRange()) { + for (auto y : variableBinAxis.GetNormalRange()) { + if (2 * x.GetIndex() == y.GetIndex()) { + EXPECT_EQ(hist->GetBinContent(x, y), 1.0); + } else { + EXPECT_EQ(hist->GetBinContent(x, y), 0.0); + } + } + } +} + +TEST_P(RDFHist, ContainerMultiDimNestedBroadcast) +{ + RDataFrame df(10); + auto dfXY = df.Define("x", [](ULong64_t e) { return e + 5.5; }, {"rdfentry_"}) + .Define("xV", [](ULong64_t e) { return ROOT::RVecD{e + 5.5}; }, {"rdfentry_"}) + .Define("y", [](ULong64_t e) { return ROOT::RVec{{2 * e + 0.5}}; }, {"rdfentry_"}); + + const RRegularAxis regularAxis(10, {5.0, 15.0}); + static constexpr std::size_t BinsY = 20; + std::vector bins; + for (std::size_t i = 0; i < BinsY; i++) { + bins.push_back(i); + } + bins.push_back(BinsY); + const RVariableBinAxis variableBinAxis(bins); + + auto hist = + dfXY.Hist>({regularAxis, variableBinAxis}, {"x", "y"}); + EXPECT_EQ(hist->GetNEntries(), 10); + for (auto x : regularAxis.GetNormalRange()) { + for (auto y : variableBinAxis.GetNormalRange()) { + if (2 * x.GetIndex() == y.GetIndex()) { + EXPECT_EQ(hist->GetBinContent(x, y), 1.0); + } else { + EXPECT_EQ(hist->GetBinContent(x, y), 0.0); + } + } + } + + hist = dfXY.Hist>({regularAxis, variableBinAxis}, + {"xV", "y"}); + EXPECT_EQ(hist->GetNEntries(), 10); +} + +TEST_P(RDFHist, ContainerInvalidElements) +{ + RDataFrame df(10); + auto dfXY = df.Define("x0", [] { return ROOT::RVecD{}; }) + .Define("x2", [](ULong64_t e) { return ROOT::RVecD{e + 5.5, e + 15.5}; }, {"rdfentry_"}) + .Define("y1", [](ULong64_t e) { return ROOT::RVecD{2 * e + 0.5}; }, {"rdfentry_"}); + + const RRegularAxis regularAxis(20, {5.0, 25.0}); + static constexpr std::size_t BinsY = 20; + std::vector bins; + for (std::size_t i = 0; i < BinsY; i++) { + bins.push_back(i); + } + bins.push_back(BinsY); + const RVariableBinAxis variableBinAxis(bins); + + try { + // Cannot use EXPECT_THROW because of template arguments... + *dfXY.Hist({regularAxis, variableBinAxis}, {"x0", "y1"}); + FAIL() << "expected std::runtime_error"; + } catch (const std::runtime_error &) { + // expected + } + + try { + // Cannot use EXPECT_THROW because of template arguments... + *dfXY.Hist({regularAxis, variableBinAxis}, {"x2", "y1"}); + FAIL() << "expected std::runtime_error"; + } catch (const std::runtime_error &e) { + // expected + } +} + +TEST_P(RDFHist, ContainerNestedInvalidElements) +{ + RDataFrame df(10); + auto dfXY = + df.Define("x11", [](ULong64_t e) { return ROOT::RVec{{e + 5.5}}; }, {"rdfentry_"}) + .Define("y12", [](ULong64_t e) { return ROOT::RVec{{2 * e + 0.5, 2 * e + 0.5}}; }, {"rdfentry_"}); + + const RRegularAxis regularAxis(20, {5.0, 25.0}); + static constexpr std::size_t BinsY = 20; + std::vector bins; + for (std::size_t i = 0; i < BinsY; i++) { + bins.push_back(i); + } + bins.push_back(BinsY); + const RVariableBinAxis variableBinAxis(bins); + + try { + // Cannot use EXPECT_THROW because of template arguments... + *dfXY.Hist, ROOT::RVec>( + {regularAxis, variableBinAxis}, {"x11", "y12"}); + FAIL() << "expected std::runtime_error"; + } catch (const std::runtime_error &e) { + // expected + } +} + +TEST_P(RDFHist, WeightContainer) +{ + RDataFrame df(10); + auto dfXW = + df.Define("xV", [](ULong64_t e) { return ROOT::RVecD{e + 5.5, e + 15.5}; }, {"rdfentry_"}) + .Define("wV", [](ULong64_t e) { return ROOT::RVecD{0.1 + e * 0.03, 0.2 + e * e * 0.06}; }, {"rdfentry_"}); + + const RRegularAxis axis(20, {5.0, 25.0}); + auto hist = dfXW.Hist({axis}, {"xV"}, "wV"); + EXPECT_EQ(hist->GetNEntries(), 20); + for (auto index : axis.GetNormalRange()) { + auto &bin = hist->GetBinContent(index); + auto i = index.GetIndex(); + double weight = i >= 10 ? (0.2 + (i - 10) * (i - 10) * 0.06) : (0.1 + i * 0.03); + EXPECT_FLOAT_EQ(bin.fSum, weight); + EXPECT_FLOAT_EQ(bin.fSum2, weight * weight); + } +} + +TEST_P(RDFHist, WeightContainerJit) +{ + RDataFrame df(10); + auto dfXW = df.Define("xV", "ROOT::RVecD{rdfentry_ + 5.5, rdfentry_ + 15.5}") + .Define("wV", "ROOT::RVecD{0.1 + rdfentry_ * 0.03, 0.2 + rdfentry_ * rdfentry_ * 0.06}"); + + const RRegularAxis axis(20, {5.0, 25.0}); + auto hist = dfXW.Hist({axis}, {"xV"}, "wV"); + EXPECT_EQ(hist->GetNEntries(), 20); + for (auto index : axis.GetNormalRange()) { + auto &bin = hist->GetBinContent(index); + auto i = index.GetIndex(); + double weight = i >= 10 ? (0.2 + (i - 10) * (i - 10) * 0.06) : (0.1 + i * 0.03); + EXPECT_FLOAT_EQ(bin.fSum, weight); + EXPECT_FLOAT_EQ(bin.fSum2, weight * weight); + } +} + +TEST_P(RDFHist, WeightContainerBroadcast) +{ + RDataFrame df(10); + auto dfXW = + df.Define("x", [](ULong64_t e) { return e + 5.5; }, {"rdfentry_"}) + .Define("xV", [](ULong64_t e) { return ROOT::RVecD{e + 5.5, e + 15.5}; }, {"rdfentry_"}) + .Define("w", [](ULong64_t e) { return 0.1 + e * 0.03; }, {"rdfentry_"}) + .Define("wV", [](ULong64_t e) { return ROOT::RVecD{0.1 + e * 0.03, 0.2 + e * e * 0.06}; }, {"rdfentry_"}); + + const RRegularAxis axis(20, {5.0, 25.0}); + auto hist = dfXW.Hist({axis}, {"xV"}, "w"); + EXPECT_EQ(hist->GetNEntries(), 20); + for (auto index : axis.GetNormalRange()) { + auto &bin = hist->GetBinContent(index); + auto i = index.GetIndex(); + double weight = 0.1 + (i >= 10 ? i - 10 : i) * 0.03; + EXPECT_FLOAT_EQ(bin.fSum, weight); + EXPECT_FLOAT_EQ(bin.fSum2, weight * weight); + } + + hist = dfXW.Hist({axis}, {"x"}, "wV"); + EXPECT_EQ(hist->GetNEntries(), 20); + EXPECT_EQ(hist->GetBinContent(2).fSum, 0.1 + 2 * 0.03 + 0.2 + 2 * 2 * 0.06); +} + INSTANTIATE_TEST_SUITE_P(Seq, RDFHist, ::testing::Values(false)); #ifdef R__USE_IMT