AliceO2Group · ktf · Dec 14, 2021 · Dec 1, 2021 · Dec 8, 2021 · Dec 8, 2021
@@ -12,6 +12,7 @@
 #ifndef FRAMEWORK_ANALYSISMANAGERS_H
 #define FRAMEWORK_ANALYSISMANAGERS_H
 #include "Framework/AnalysisHelpers.h"
+#include "Framework/GroupedCombinations.h"
 #include "Framework/Kernels.h"
 #include "Framework/ASoA.h"
 #include "Framework/ProcessingContext.h"
@@ -29,6 +30,29 @@
 namespace o2::framework
 {

+template <typename ANY>
+struct GroupedCombinationManager {
+  template <typename TG, typename... T2s>
+  static void setGroupedCombination(ANY&, TG&, T2s&...)
+  {
+  }
+};
+
+template <typename T1, typename GroupingPolicy, typename H, typename G, typename... Us, typename... As>
+struct GroupedCombinationManager<GroupedCombinationsGenerator<T1, GroupingPolicy, H, G, pack<Us...>, As...>> {
+  template <typename TH, typename TG, typename... T2s>
+  static void setGroupedCombination(GroupedCombinationsGenerator<T1, GroupingPolicy, H, G, pack<Us...>, As...>& comb, TH& hashes, TG& grouping, std::tuple<T2s...>& associated)
+  {
+    static_assert(sizeof...(T2s) > 0, "There must be associated tables in process() for a correct pair");
+    static_assert(!soa::is_soa_iterator_t<std::decay_t<H>>::value, "Only full tables can be in process(), no grouping");
+    if constexpr (std::conjunction_v<std::is_same<G, TG>, std::is_same<H, TH>>) {
+      // Take respective unique associated tables for grouping
+      auto associatedTuple = std::tuple<Us...>(std::get<Us>(associated)...);
+      comb.setTables(hashes, grouping, associatedTuple);
+    }
+  }
+};
+
 template <typename ANY>
 struct PartitionManager {
  template <typename... T2s>

@@ -55,7 +55,6 @@ struct AnalysisTask {

 // Helper struct which builds a DataProcessorSpec from
 // the contents of an AnalysisTask...
-
 struct AnalysisDataProcessorBuilder {
  template <typename T>
  static ConfigParamSpec getSpec()
@@ -262,6 +261,7 @@ struct AnalysisDataProcessorBuilder {
      // single argument to process
      homogeneous_apply_refs([&groupingTable](auto& x) {
        PartitionManager<std::decay_t<decltype(x)>>::bindExternalIndices(x, &groupingTable);
+        GroupedCombinationManager<std::decay_t<decltype(x)>>::setGroupedCombination(x, groupingTable);
        return true;
      },
                             task);
@@ -310,6 +310,15 @@ struct AnalysisDataProcessorBuilder {
        },
        associatedTables);

+      // GroupedCombinations bound separately, as they should be set once for all associated tables
+      auto hashes = std::get<0>(associatedTables);
+      auto realAssociated = tuple_tail(associatedTables);
+      homogeneous_apply_refs([&groupingTable, &hashes, &realAssociated](auto& t) {
+        GroupedCombinationManager<std::decay_t<decltype(t)>>::setGroupedCombination(t, hashes, groupingTable, realAssociated);
+        return true;
+      },
+                             task);
+
      if constexpr (soa::is_soa_iterator_t<std::decay_t<G>>::value) {
        // grouping case
        auto slicer = GroupSlicer(groupingTable, associatedTables);

@@ -0,0 +1,267 @@
+// Copyright 2019-2020 CERN and copyright holders of ALICE O2.
+// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
+// All rights not expressly granted are reserved.
+//
+// This software is distributed under the terms of the GNU General Public
+// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
+//
+// In applying this license CERN does not waive the privileges and immunities
+// granted to it by virtue of its status as an Intergovernmental Organization
+// or submit itself to any jurisdiction.
+
+#ifndef FRAMEWORK_GROUPEDCOMBINATIONS_H
+#define FRAMEWORK_GROUPEDCOMBINATIONS_H
+
+#include "Framework/ASoAHelpers.h"
+#include "Framework/GroupSlicer.h"
+#include "Framework/Pack.h"
+
+namespace o2::framework
+{
+
+// Create an instance of a tuple interleaved from given tuples
+template <typename... T1s, typename... T2s, std::size_t... Is>
+auto interleaveTuplesImpl(std::tuple<T1s...>& t1, std::tuple<T2s...>& t2, std::index_sequence<Is...>)
+{
+  return std::tuple_cat(std::make_tuple(std::get<Is>(t1), std::get<Is>(t2))...);
+}
+
+template <typename... T1s, typename... T2s>
+auto interleaveTuples(std::tuple<T1s...>& t1, std::tuple<T2s...>& t2)
+{
+  return interleaveTuplesImpl(t1, t2, std::index_sequence_for<T1s...>());
+}
+
+// Functions to create a tuple from N runs of a function that returns a value
+template <std::size_t I, typename R, typename C, typename... Args>
+R execFunctionWithDummyIndex(R (C::*f)(Args...), C& obj, Args... args)
+{
+  return (obj.*f)(args...);
+}
+
+template <typename R, typename C, typename... Args, std::size_t... Is>
+auto functionToTupleImpl(R (C::*f)(Args...), C& obj, Args... args, std::index_sequence<Is...>)
+{
+  return std::make_tuple((execFunctionWithDummyIndex<Is>(f, obj, args...))...);
+}
+
+template <std::size_t N, typename R, typename C, typename... Args>
+auto functionToTuple(R (C::*f)(Args...), C& obj, Args... args)
+{
+  return functionToTupleImpl(f, obj, args..., std::make_index_sequence<N>());
+}
+
+template <typename T1, typename GroupingPolicy, typename H, typename G, typename... Ts>
+struct GroupedCombinationsGenerator {
+};
+
+template <typename T1, typename GroupingPolicy, typename H, typename G, typename... Us, typename... As>
+struct GroupedCombinationsGenerator<T1, GroupingPolicy, H, G, pack<Us...>, As...> {
+  using joinIterator = typename soa::Join<H, G>::table_t::iterator;
+  using GroupedIteratorType = pack_to_tuple_t<interleaved_pack_t<repeated_type_pack_t<joinIterator, sizeof...(As)>, pack<As...>>>;
+
+  struct GroupedIterator : public std::iterator<std::forward_iterator_tag, GroupedIteratorType>, public GroupingPolicy {
+   public:
+    using reference = GroupedIteratorType&;
+    using value_type = GroupedIteratorType;
+    using pointer = GroupedIteratorType*;
+    using iterator_category = std::forward_iterator_tag;
+
+    GroupedIterator(const GroupingPolicy& groupingPolicy) : GroupingPolicy(groupingPolicy) {}
+    GroupedIterator(const GroupingPolicy& groupingPolicy, const H& hashes, const G& grouping, const std::shared_ptr<GroupSlicer<G, Us...>>&& slicer_ptr) : GroupingPolicy(groupingPolicy), mSlicer{std::move(slicer_ptr)}, mGrouping{std::make_shared<G>(std::vector{grouping.asArrowTable()})}
+    {
+      GroupingPolicy::setTables(join(hashes, grouping), join(hashes, grouping));
+      if (!this->mIsEnd) {
+        setCurrentGroupedCombination();
+      }
+    }
+
+    GroupedIterator(GroupedIterator const&) = default;
+    GroupedIterator& operator=(GroupedIterator const&) = default;
+    ~GroupedIterator() = default;
+
+    void setTables(const H& hashes, const G& grouping, std::shared_ptr<GroupSlicer<G, Us...>> slicer_ptr)
+    {
+      mGrouping = std::make_shared<G>(std::vector{grouping.asArrowTable()});
+      mSlicer = slicer_ptr;
+      setMultipleGroupingTables<sizeof...(As)>(join(hashes, grouping));
+      if (!this->mIsEnd) {
+        setCurrentGroupedCombination();
+      }
+    }
+
+    template <std::size_t N, typename T, typename... Args>
+    void setMultipleGroupingTables(const T& param, const Args&... args)
+    {
+      if constexpr (N == 1) {
+        GroupingPolicy::setTables(param, args...);
+      } else {
+        setMultipleGroupingTables<N - 1>(param, param, args...);
+      }
+    }
+
+    void moveToEnd()
+    {
+      GroupingPolicy::moveToEnd();
+    }
+
+    // prefix increment
+    GroupedIterator& operator++()
+    {
+      if (!this->mIsEnd) {
+        this->addOne();
+        setCurrentGroupedCombination();
+      }
+      return *this;
+    }
+    // postfix increment
+    GroupedIterator operator++(int /*unused*/)
+    {
+      GroupedIterator copy(*this);
+      operator++();
+      return copy;
+    }
+    // return reference
+    reference operator*()
+    {
+      return *mCurrentGrouped;
+    }
+    bool operator==(const GroupedIterator& rh)
+    {
+      return (this->mIsEnd && rh.mIsEnd) || (this->mCurrent == rh.mCurrent);
+    }
+    bool operator!=(const GroupedIterator& rh)
+    {
+      return !(*this == rh);
+    }
+
+   private:
+    std::tuple<As...> getAssociatedTables()
+    {
+      auto& currentGrouping = GroupingPolicy::mCurrent;
+      constexpr auto k = sizeof...(As);
+      auto slicerIterators = functionToTuple<k>(&GroupSlicer<G, Us...>::begin, *mSlicer);
+      o2::soa::for_<k>([&](auto i) {
+        auto col = std::get<i.value>(currentGrouping);
+        for (auto& slice : *mSlicer) {
+          if (slice.groupingElement().globalIndex() == col.globalIndex()) {
+            std::get<i.value>(slicerIterators) = slice;
+            break;
+          }
+        }
+      });
+
+      return getSlices(slicerIterators, std::index_sequence_for<As...>());
+    }
+
+    template <std::size_t I, typename T, typename... Ts>
+    auto getSliceAt(std::tuple<Ts...>& t)
+    {
+      auto it = std::get<I>(t); // Get the tables corresponding to the grouping at index I
+      auto associatedType = it.template prepareArgument<T>();
+      return associatedType;
+    }
+
+    template <typename... Ts, std::size_t... Is>
+    std::tuple<As...> getSlices(std::tuple<Ts...>& t, std::index_sequence<Is...> is)
+    {
+      return std::make_tuple(getSliceAt<Is, As, Ts...>(t)...);
+    }
+
+    void setCurrentGroupedCombination()
+    {
+      std::tuple<As...> initAssociatedTables = getAssociatedTables();
+      constexpr auto k = sizeof...(As);
+      bool moveForward = false;
+      o2::soa::for_<k>([&](auto i) {
+        if (std::get<i.value>(initAssociatedTables).size() == 0) {
+          moveForward = true;
+        }
+      });
+      while (!this->mIsEnd && moveForward) {
+        GroupingPolicy::addOne();
+        std::tuple<As...> temp = getAssociatedTables();
+        moveForward = false;
+        o2::soa::for_<k>([&](auto i) {
+          if (std::get<i.value>(temp).size() == 0) {
+            moveForward = true;
+          }
+        });
+      }
+      std::tuple<As...> associatedTables = getAssociatedTables();
+
+      if (!this->mIsEnd) {
+        auto& currentGrouping = GroupingPolicy::mCurrent;
+        o2::soa::for_<k>([&](auto i) {
+          std::get<i.value>(associatedTables).bindExternalIndices(mGrouping.get());
+        });
+
+        mCurrentGrouped.emplace(interleaveTuples(currentGrouping, associatedTables));
+      }
+    }
+
+    std::shared_ptr<GroupSlicer<G, Us...>> mSlicer = nullptr;
+    std::shared_ptr<G> mGrouping;
+    std::optional<GroupedIteratorType> mCurrentGrouped;
+  };
+
+  using iterator = GroupedIterator;
+  using const_iterator = GroupedIterator;
+
+  inline iterator begin()
+  {
+    return iterator(mBegin);
+  }
+  inline iterator end()
+  {
+    return iterator(mEnd);
+  }
+  inline const_iterator begin() const
+  {
+    return iterator(mBegin);
+  }
+  inline const_iterator end() const
+  {
+    return iterator(mEnd);
+  }
+
+  GroupedCombinationsGenerator(const char* category, int catNeighbours, const T1& outsider) : mBegin(GroupingPolicy(category, catNeighbours, outsider)), mEnd(GroupingPolicy(category, catNeighbours, outsider)), mCategory(category), mCatNeighbours(catNeighbours), mOutsider(outsider) {}
+  GroupedCombinationsGenerator(const char* category, int catNeighbours, const T1& outsider, H& hashes, G& grouping, std::tuple<Us...>& associated) : GroupedCombinationsGenerator(category, catNeighbours, outsider)
+  {
+    setTables(hashes, grouping, associated);
+  }
+  ~GroupedCombinationsGenerator() = default;
+
+  void setTables(H& hashes, G& grouping, std::tuple<Us...>& associated)
+  {
+    std::shared_ptr slicer_ptr = std::make_shared<GroupSlicer<G, Us...>>(grouping, associated);
+    mBegin.setTables(hashes, grouping, slicer_ptr);
+    mEnd.setTables(hashes, grouping, slicer_ptr);
+    mEnd.moveToEnd();
+  }
+
+ private:
+  iterator mBegin;
+  iterator mEnd;
+  const char* mCategory;
+  const int mCatNeighbours;
+  const T1 mOutsider;
+};
+
+// Aliases for 2-particle correlations
+// 'Pair' and 'Triple' can be used for same kind pair/triple, too, just specify the same type twice
+template <typename H, typename G>
+using joinedCollisions = typename soa::Join<H, G>::table_t;
+template <typename H, typename G, typename A1, typename A2, typename T1 = int, typename GroupingPolicy = o2::soa::CombinationsBlockStrictlyUpperSameIndexPolicy<T1, joinedCollisions<H, G>, joinedCollisions<H, G>>>
+using Pair = GroupedCombinationsGenerator<T1, GroupingPolicy, H, G, unique_pack_t<pack<A1, A2>>>;
+template <typename H, typename G, typename A, typename T1 = int, typename GroupingPolicy = o2::soa::CombinationsBlockStrictlyUpperSameIndexPolicy<T1, joinedCollisions<H, G>, joinedCollisions<H, G>>>
+using SameKindPair = GroupedCombinationsGenerator<T1, GroupingPolicy, H, G, pack<A>, A, A>;
+
+// Aliases for 3-particle correlations
+template <typename H, typename G, typename A1, typename A2, typename A3, typename T1 = int, typename GroupingPolicy = o2::soa::CombinationsBlockStrictlyUpperSameIndexPolicy<T1, joinedCollisions<H, G>, joinedCollisions<H, G>, joinedCollisions<H, G>>>
+using Triple = GroupedCombinationsGenerator<T1, GroupingPolicy, H, G, unique_pack_t<pack<A1, A2, A3>>>;
+template <typename H, typename G, typename A, typename T1 = int, typename GroupingPolicy = o2::soa::CombinationsBlockStrictlyUpperSameIndexPolicy<T1, joinedCollisions<H, G>, joinedCollisions<H, G>, joinedCollisions<H, G>>>
+using SameKindTriple = GroupedCombinationsGenerator<T1, GroupingPolicy, H, G, pack<A>, A, A, A>;
+
+} // namespace o2::framework
+#endif // FRAMEWORK_GROUPEDCOMBINATIONS_H_
@@ -272,6 +272,24 @@ constexpr auto unique_pack(pack<T, Ts...>, PT p2)
 template <typename P>
 using unique_pack_t = decltype(unique_pack(P{}, pack<>{}));

+template <typename... Ts>
+inline constexpr std::tuple<Ts...> pack_to_tuple(pack<Ts...>)
+{
+  return std::tuple<Ts...>{};
+}
+
+template <typename P>
+using pack_to_tuple_t = decltype(pack_to_tuple(P{}));
+
+template <typename T, std::size_t... Is>
+inline auto sequence_to_pack(std::integer_sequence<std::size_t, Is...>)
+{
+  return pack<decltype((Is, T{}))...>{};
+};
+
+template <typename T, std::size_t N>
+using repeated_type_pack_t = decltype(sequence_to_pack<T>(std::make_index_sequence<N>()));
+
 } // namespace o2::framework

 #endif // O2_FRAMEWORK_PACK_H_
@@ -53,6 +53,8 @@ BOOST_AUTO_TEST_CASE(TestOverride)

  static_assert(std::is_same_v<unique_pack_t<pack<int, float, int, float, char, char>>, pack<char, float, int>>, "pack should not have duplicated types");
  static_assert(std::is_same_v<interleaved_pack_t<pack<int, float, int>, pack<char, bool, char>>, pack<int, char, float, bool, int, char>>, "interleaved packs of the same size");
+  static_assert(std::is_same_v<pack_to_tuple_t<pack<int, float, char>>, std::tuple<int, float, char>>, "pack should become a tuple");
+  static_assert(std::is_same_v<repeated_type_pack_t<float, 5>, pack<float, float, float, float, float>>, "pack should have float repeated 5 times");

  struct ForwardDeclared;
  static_assert(is_type_complete_v<ForwardDeclared> == false, "This should not be complete because the struct is simply forward declared.");