From eaf1d69a99269d4e3ae6e62f38df67bb1728a8f9 Mon Sep 17 00:00:00 2001
From: Giulio Eulisse <10544+ktf@users.noreply.github.com>
Date: Tue, 26 May 2026 21:04:33 +0200
Subject: [PATCH] StepTHn: speedup filling

Replaces 4 virtual calls to GetAt / SetAt / AddAt  with just one to updateBin.
---
 Framework/Core/include/Framework/StepTHn.h | 25 +++++++++++++++++++++-
 Framework/Core/src/StepTHn.cxx             | 21 ++----------------
 2 files changed, 26 insertions(+), 20 deletions(-)

diff --git a/Framework/Core/include/Framework/StepTHn.h b/Framework/Core/include/Framework/StepTHn.h
index 0302f604eae39..165ad8270f206 100644
--- a/Framework/Core/include/Framework/StepTHn.h
+++ b/Framework/Core/include/Framework/StepTHn.h
@@ -58,7 +58,7 @@ class StepTHn : public TNamed
   virtual Long64_t Merge(TCollection* list) = 0;
 
   TAxis* GetAxis(int i) { return mPrototype->GetAxis(i); }
-  void Sumw2(){}; // TODO: added for compatibiltiy with registry, but maybe it would be useful also in StepTHn as toggle for error weights
+  void Sumw2() {}; // TODO: added for compatibiltiy with registry, but maybe it would be useful also in StepTHn as toggle for error weights
 
  protected:
   void init();
@@ -67,6 +67,7 @@ class StepTHn : public TNamed
   void deleteContainers();
 
   Long64_t getGlobalBinIndex(const Int_t* binIdx);
+  virtual void updateBin(int iStep, Long64_t bin, double weight) = 0;
 
   Long64_t mNBins;  // number of total bins
   Int_t mNVars;     // number of variables
@@ -107,6 +108,28 @@ class StepTHnT : public StepTHn
     }
   }
 
+  void updateBin(int iStep, Long64_t bin, double weight) override
+  {
+    if (!mValues[iStep]) {
+      mValues[iStep] = createArray();
+      LOGF(info, "Created values container for step %d", iStep);
+    }
+
+    if (weight != 1.) {
+      if (!mSumw2[iStep]) {
+        mSumw2[iStep] = createArray(mValues[iStep]);
+        LOGF(info, "Created sumw2 container for step %d", iStep);
+      }
+    }
+
+    auto* arr = static_cast<TemplateArray*>(mValues[iStep])->GetArray();
+    arr[bin] += weight;
+    if (mSumw2[iStep]) {
+      auto* sw2 = static_cast<TemplateArray*>(mSumw2[iStep])->GetArray();
+      sw2[bin] += weight * weight;
+    }
+  }
+
   ClassDef(StepTHnT, 1) // THn like container
 };
 
diff --git a/Framework/Core/src/StepTHn.cxx b/Framework/Core/src/StepTHn.cxx
index bb0109db2c97f..ed4a544a44c44 100644
--- a/Framework/Core/src/StepTHn.cxx
+++ b/Framework/Core/src/StepTHn.cxx
@@ -424,7 +424,7 @@ void StepTHn::Fill(int iStep, int nParams, double positionAndWeight[])
       mLastBins[i] = tmpBin;
       mLastVars[i] = positionAndWeight[i];
     }
-    //Printf("%d", tmpBin);
+    // Printf("%d", tmpBin);
 
     // under/overflow not supported
     if (tmpBin < 1 || tmpBin > mNbinsCache[i]) {
@@ -436,24 +436,7 @@ void StepTHn::Fill(int iStep, int nParams, double positionAndWeight[])
     //     Printf("%lld", bin);
   }
 
-  if (!mValues[iStep]) {
-    mValues[iStep] = createArray();
-    LOGF(info, "Created values container for step %d", iStep);
-  }
-
-  if (weight != 1.) {
-    // initialize with already filled entries (which have been filled with weight == 1), in this case mSumw2 := mValues
-    if (!mSumw2[iStep]) {
-      mSumw2[iStep] = createArray(mValues[iStep]);
-      LOGF(info, "Created sumw2 container for step %d", iStep);
-    }
-  }
-
-  // TODO probably slow; add StepTHnT::add ?
-  mValues[iStep]->SetAt(mValues[iStep]->GetAt(bin) + weight, bin);
-  if (mSumw2[iStep]) {
-    mSumw2[iStep]->SetAt(mSumw2[iStep]->GetAt(bin) + weight * weight, bin);
-  }
+  updateBin(iStep, bin, weight);
 }
 
 template class StepTHnT<TArrayF>;