add occupancy-related structures and enums to EventSelection

Author: altsybee

Common/DataModel/EventSelection.h

@@ -18,6 +18,12 @@
 #define COMMON_DATAMODEL_EVENTSELECTION_H_
 
 #include <Framework/AnalysisDataModel.h>
+#include <Framework/Logger.h>
+
+#include <algorithm>
+#include <cstdint>
+#include <ranges>
+#include <vector>
 
 namespace o2::aod
 {
@@ -85,6 +91,105 @@ DECLARE_SOA_TABLE(EvSels, "AOD", "EVSEL", //!
                   evsel::SumAmpFT0CInTimeRange,
                   evsel::OccupancyMedianTime);
 using EvSel = EvSels::iterator;
+
+// occupancy wrt BCs: same estimators and flags as in EvSels (occupancy from other collisions nearby in time),
+// calculated wrt a BC; collisions found in this BC are excluded (their number is stored in numCollsInBC)
+// N.B.: not related to the occupancy tables in OccupancyTables.h
+namespace bcocc
+{
+enum BcOccupancyFlags {         // bits of occupancyFlags
+  kNoCollInTimeRangeNarrow = 0, // no other collisions in specified time range (narrower than Strict)
+  kNoCollInTimeRangeStrict,     // no other collisions in specified time range
+  kNoCollInTimeRangeStandard,   // no other collisions in specified time range with per-collision multiplicity above threshold
+  kNoCollInRofStrict,           // no other collisions in this ITS Readout Frame
+  kNoCollInRofStandard,         // no other collisions in this ITS Readout Frame with per-collision multiplicity above threshold (no veto on collisions with close vZ)
+  kNoHighMultCollInPrevRof      // FT0C amplitude in previous ITS ROF is below threshold
+};
+DECLARE_SOA_COLUMN(NumTracksInTimeRange, trackOccupancyInTimeRange, int);   //! Occupancy in specified time interval by a number of tracks from nearby collisions (-1 if not available) // o2-linter: disable=name/o2-column
+DECLARE_SOA_COLUMN(SumAmpFT0CInTimeRange, ft0cOccupancyInTimeRange, float); //! Occupancy in specified time interval by a sum of FT0C amplitudes from nearby collisions (-1 if not available) // o2-linter: disable=name/o2-column
+DECLARE_SOA_COLUMN(OccupancyMedianTime, occupancyMedianTime, float);        //! Median time for the occupancy in specified time interval (the delta-time wrt to the BC at which occupancy crosses 1/2 of its full value)
+DECLARE_SOA_BITMAP_COLUMN(OccupancyFlags, occupancyFlags, 8);               //! Bitmask of occupancy flags (see BcOccupancyFlags), check with occupancyFlags_bit(bit)
+DECLARE_SOA_COLUMN(NumCollsInBC, numCollsInBC, int16_t);                    //! Number of collisions found in this BC (excluded from the occupancy)
+} // namespace bcocc
+
+// bc-joinable occupancy
+DECLARE_SOA_TABLE(BcOccs, "AOD", "BCOCC", //!
+                  bcocc::NumTracksInTimeRange,
+                  bcocc::SumAmpFT0CInTimeRange,
+                  bcocc::OccupancyMedianTime,
+                  bcocc::OccupancyFlags,
+                  bcocc::NumCollsInBC);
+using BcOcc = BcOccs::iterator;
+
+// occupancy wrt every BC of the TFs present in the DF (not joinable): for each TF (in increasing order) all its BCs, in increasing order of global BC
+// same estimators and flags as in BcOccs
+DECLARE_SOA_TABLE(BcOccsAll, "AOD", "BCOCCALL", //!
+                  bc::GlobalBC,
+                  bcocc::NumTracksInTimeRange,
+                  bcocc::SumAmpFT0CInTimeRange,
+                  bcocc::OccupancyMedianTime,
+                  bcocc::OccupancyFlags,
+                  bcocc::NumCollsInBC);
+using BcOccAll = BcOccsAll::iterator;
+
+namespace bcocc
+{
+// fast lookup of the aod::BcOccsAll entry for a given global BC
+// usage (within the process function, the index keeps a pointer to the table):
+//   aod::bcocc::BcOccAllIndex bcOccAllIndexed(bcOccsAll);
+//   auto occ = bcOccAllIndexed[bc.globalBC()]; // iterator of aod::BcOccsAll
+//   int occByTracks = occ.trackOccupancyInTimeRange();
+template <typename TBcOccsAll>
+class BcOccAllIndex
+{
+ public:
+  explicit BcOccAllIndex(TBcOccsAll const& bcOccsAll) : mTable(&bcOccsAll)
+  {
+    const int64_t nRows = static_cast<int64_t>(bcOccsAll.size());
+    auto offsetAt = [&bcOccsAll](int64_t row) { return static_cast<int64_t>(bcOccsAll.iteratorAt(row).globalBC()) - row; };
+    for (int64_t first = 0, end = 0; first < nRows; first = end) {
+      // the offset is non-decreasing (global BCs are strictly increasing): the segment ends at the first row with a larger offset
+      const int64_t offset = offsetAt(first);
+      end = *std::ranges::partition_point(std::views::iota(first, nRows), [&](int64_t row) { return offsetAt(row) == offset; });
+      mSegments.push_back({offset + first, offset + end, offset});
+    }
+  }
+
+  // entry of aod::BcOccsAll for this global BC (fatal if the BC is not in the table)
+  auto operator[](uint64_t globalBC) const
+  {
+    const int64_t row = rowOf(globalBC);
+    if (row < 0) {
+      LOGP(fatal, "BcOccAllIndex: global BC {} is not in aod::BcOccsAll ({} rows): its TF is not in this DF, or the occupancy is not available for this DF", globalBC, mTable->size());
+    }
+    return mTable->iteratorAt(row);
+  }
+
+  // helper function to check if a global BC is contained in the table
+  bool contains(uint64_t globalBC) const { return rowOf(globalBC) >= 0; }
+
+ private:
+  // row of aod::BcOccsAll for this global BC, -1 if the BC is not in the table
+  int64_t rowOf(uint64_t globalBC) const
+  {
+    const auto bc = static_cast<int64_t>(globalBC);
+    for (const auto& seg : mSegments) { // a few segments: one per range of consecutive TFs in the DF
+      if (bc >= seg.firstGlobalBC && bc < seg.endGlobalBC) {
+        return bc - seg.offset;
+      }
+    }
+    return -1;
+  }
+
+  struct Segment {
+    int64_t firstGlobalBC; // first global BC of the segment
+    int64_t endGlobalBC;   // last global BC of the segment + 1
+    int64_t offset;        // globalBC - row, constant within the segment
+  };
+  TBcOccsAll const* mTable;       // not owned: the index is valid as long as the table (within the process function)
+  std::vector<Segment> mSegments; // in increasing order of global BC
+};
+} // namespace bcocc
 } // namespace o2::aod
 
 #endif // COMMON_DATAMODEL_EVENTSELECTION_H_