StepTHn: do a pre lookup for the bin using a uniform grid

Rather than doing a O(log N) binary search for the correct bin, precompute a uniform index for a O(1) lookup and find the actual bin using a short linear search.

Author: ktf

Framework/Core/include/Framework/StepTHn.h

@@ -77,10 +77,22 @@ class StepTHn : public TNamed
 
   THnBase** mTarget; //! target histogram
 
-  TAxis** mAxisCache;  //! cache axis pointers (about 50% of the time in Fill is spent in GetAxis otherwise)
-  Int_t* mNbinsCache;  //! cache Nbins per axis
-  Double_t* mLastVars; //! caching of last used bins (in many loops some vars are the same for a while)
-  Int_t* mLastBins;    //! caching of last used bins (in many loops some vars are the same for a while)
+  TAxis** mAxisCache;   //! cache axis pointers (about 50% of the time in Fill is spent in GetAxis otherwise)
+  Int_t* mNbinsCache;   //! cache Nbins per axis
+  Double_t* mLastVars;  //! caching of last used bins (in many loops some vars are the same for a while)
+  Int_t* mLastBins;     //! caching of last used bins (in many loops some vars are the same for a while)
+
+  // Fast bin lookup table: for each axis, maps a quantized position to an approximate bin.
+  static constexpr Int_t kLookupSize = 1024;  // number of slots per axis
+  struct AxisLookup {
+    Double_t invSlotWidth; // 1.0 / slot width for fast index computation
+    Double_t xmin;         // axis minimum
+    Double_t xmax;         // axis maximum
+    const Double_t* edges; // pointer to bin edges array (nBins+1 entries)
+    Int_t nBins;           // number of bins
+    Int_t table[kLookupSize]; // slot -> bin index (1-based, TAxis convention)
+  };
+  AxisLookup* mLookup; //! per-axis lookup tables
 
   THnSparse* mPrototype; // not filled used as prototype histogram for axis functionality etc.
 

Framework/Core/src/StepTHn.cxx

@@ -39,6 +39,7 @@ StepTHn::StepTHn() : mNBins(0),
                      mNbinsCache(nullptr),
                      mLastVars(nullptr),
                      mLastBins(nullptr),
+                     mLookup(nullptr),
                      mPrototype(nullptr)
 {
   // Default constructor (for streaming)
@@ -55,6 +56,7 @@ StepTHn::StepTHn(const Char_t* name, const Char_t* title, const Int_t nSteps, co
                                                                                                    mNbinsCache(nullptr),
                                                                                                    mLastVars(nullptr),
                                                                                                    mLastBins(nullptr),
+                                                                                                   mLookup(nullptr),
                                                                                                    mPrototype(nullptr)
 {
   // Constructor
@@ -130,6 +132,7 @@ StepTHn::StepTHn(const StepTHn& c) : mNBins(c.mNBins),
                                      mNbinsCache(nullptr),
                                      mLastVars(nullptr),
                                      mLastBins(nullptr),
+                                     mLookup(nullptr),
                                      mPrototype(nullptr)
 {
   //
@@ -152,6 +155,7 @@ StepTHn::~StepTHn()
   delete[] mNbinsCache;
   delete[] mLastVars;
   delete[] mLastBins;
+  delete[] mLookup;
   delete mPrototype;
 }
 
@@ -392,13 +396,35 @@ void StepTHn::Fill(int iStep, int nParams, double positionAndWeight[])
     LOGF(fatal, "Fill called with invalid number of parameters (%d vs %d)", mNVars, nParams);
   }
 
-  // fill axis cache
+  // fill axis cache and build lookup tables on first call
   if (!mAxisCache) {
     mAxisCache = new TAxis*[mNVars];
     mNbinsCache = new Int_t[mNVars];
+    mLookup = new AxisLookup[mNVars];
     for (Int_t i = 0; i < mNVars; i++) {
       mAxisCache[i] = mPrototype->GetAxis(i);
       mNbinsCache[i] = mAxisCache[i]->GetNbins();
+
+      // Build lookup table for this axis
+      auto& lut = mLookup[i];
+      lut.nBins = mNbinsCache[i];
+      lut.edges = mAxisCache[i]->GetXbins()->GetArray();
+      lut.xmin = mAxisCache[i]->GetXmin();
+      lut.xmax = mAxisCache[i]->GetXmax();
+
+      if (lut.edges && lut.xmax > lut.xmin) {
+        // Variable-width binning: build slot->bin mapping
+        Double_t slotWidth = (lut.xmax - lut.xmin) / kLookupSize;
+        lut.invSlotWidth = 1.0 / slotWidth;
+        for (Int_t s = 0; s < kLookupSize; s++) {
+          Double_t x = lut.xmin + (s + 0.5) * slotWidth;
+          lut.table[s] = mAxisCache[i]->FindBin(x);
+        }
+      } else {
+        // Uniform binning or degenerate axis: direct formula, no table needed
+        lut.edges = nullptr;
+        lut.invSlotWidth = lut.xmax > lut.xmin ? lut.nBins / (lut.xmax - lut.xmin) : 0;
+      }
     }
 
     mLastVars = new Double_t[mNVars];
@@ -420,11 +446,39 @@ void StepTHn::Fill(int iStep, int nParams, double positionAndWeight[])
     if (mLastVars[i] == positionAndWeight[i]) {
       tmpBin = mLastBins[i];
     } else {
-      tmpBin = mAxisCache[i]->FindBin(positionAndWeight[i]);
+      const auto& lut = mLookup[i];
+      Double_t x = positionAndWeight[i];
+
+      if (!lut.edges) {
+        // Uniform binning: direct computation
+        tmpBin = 1 + static_cast<Int_t>((x - lut.xmin) * lut.invSlotWidth);
+        if (tmpBin < 1) {
+          tmpBin = 0; // underflow
+        } else if (tmpBin > lut.nBins) {
+          tmpBin = lut.nBins + 1; // overflow
+        }
+      } else {
+        // Variable-width binning: lookup table + refine
+        Int_t slot = static_cast<Int_t>((x - lut.xmin) * lut.invSlotWidth);
+        if (slot < 0 || slot >= kLookupSize) {
+          tmpBin = (slot < 0) ? 0 : lut.nBins + 1; // under/overflow
+        } else {
+          tmpBin = lut.table[slot];
+          // Refine: the lookup gives an approximate bin, check edges
+          // Move left if x is below the lower edge of tmpBin
+          while (tmpBin > 1 && x < lut.edges[tmpBin - 1]) {
+            tmpBin--;
+          }
+          // Move right if x is at or above the upper edge of tmpBin
+          while (tmpBin < lut.nBins && x >= lut.edges[tmpBin]) {
+            tmpBin++;
+          }
+        }
+      }
+
       mLastBins[i] = tmpBin;
-      mLastVars[i] = positionAndWeight[i];
+      mLastVars[i] = x;
     }
-    // Printf("%d", tmpBin);
 
     // under/overflow not supported
     if (tmpBin < 1 || tmpBin > mNbinsCache[i]) {
@@ -433,7 +487,6 @@ void StepTHn::Fill(int iStep, int nParams, double positionAndWeight[])
 
     // bins start from 0 here
     bin += tmpBin - 1;
-    //     Printf("%lld", bin);
   }
 
   updateBin(iStep, bin, weight);