From a25d9c42c0c68f2735fd3c63d578da7b7671c308 Mon Sep 17 00:00:00 2001
From: Francesco Mazzaschi <fmazzasc@alipap1.cern.ch>
Date: Thu, 24 Sep 2026 10:27:55 +0200
Subject: [PATCH] BoxGenerator: enable sampling of pT and rapidity instead of p
 and eta

---
 Generators/include/Generators/BoxGenerator.h | 39 ++++++++++++++++----
 Generators/include/Generators/BoxGunParam.h  |  5 ++-
 Generators/src/BoxGenerator.cxx              |  7 +++-
 3 files changed, 41 insertions(+), 10 deletions(-)

diff --git a/Generators/include/Generators/BoxGenerator.h b/Generators/include/Generators/BoxGenerator.h
index e109bcf90ebf8..ff0b9cbfcf75c 100644
--- a/Generators/include/Generators/BoxGenerator.h
+++ b/Generators/include/Generators/BoxGenerator.h
@@ -35,6 +35,7 @@ class BoxGenerator : public Generator
   BoxGenerator() = default;
   BoxGenerator(int pdgid, int mult = 1);
 
+  /// With sampleYAndPt, eta bounds specify rapidity and p bounds specify pT.
   BoxGenerator(int pdgid,
                int mult,
                double etamin,
@@ -42,18 +43,23 @@ class BoxGenerator : public Generator
                double pmin,
                double pmax,
                double phimin,
-               double phimax) : mPDG{pdgid}, mMult{mult}
+               double phimax,
+               bool sampleYAndPt = false) : mPDG{pdgid}, mMult{mult}
   {
-    SetEtaRange(etamin, etamax);
-    SetPRange(pmin, pmax);
+    if (sampleYAndPt) {
+      SetYRange(etamin, etamax);
+      SetPtRange(pmin, pmax);
+    } else {
+      SetEtaRange(etamin, etamax);
+      SetPRange(pmin, pmax);
+    }
     SetPhiRange(phimin, phimax);
   }
 
-  BoxGenerator(BoxGenConfig const& config) : mPDG{config.pdg}, mMult{config.number}
+  BoxGenerator(BoxGenConfig const& config)
+    : BoxGenerator(config.pdg, config.number, config.eta[0], config.eta[1],
+                   config.prange[0], config.prange[1], config.phirange[0], config.phirange[1], config.sampleYAndPt)
   {
-    SetEtaRange(config.eta[0], config.eta[1]);
-    SetPRange(config.prange[0], config.prange[1]);
-    SetPhiRange(config.phirange[0], config.phirange[1]);
   }
 
   void SetPRange(Double32_t pmin = 0, Double32_t pmax = 10)
@@ -61,6 +67,15 @@ class BoxGenerator : public Generator
     mPMin = pmin;
     mPMax = pmax;
     mPRangeIsSet = true;
+    mPtRangeIsSet = false;
+  }
+
+  void SetPtRange(Double32_t ptmin = 0, Double32_t ptmax = 10)
+  {
+    mPtMin = ptmin;
+    mPtMax = ptmax;
+    mPtRangeIsSet = true;
+    mPRangeIsSet = false;
   }
 
   void SetPhiRange(double phimin = 0, double phimax = 360)
@@ -74,6 +89,16 @@ class BoxGenerator : public Generator
     mEtaMin = etamin;
     mEtaMax = etamax;
     mEtaRangeIsSet = true;
+    mYRangeIsSet = false;
+  }
+
+  /// Sample rapidity uniformly; requires a transverse momentum range.
+  void SetYRange(double ymin = -5, double ymax = 5)
+  {
+    mYMin = ymin;
+    mYMax = ymax;
+    mYRangeIsSet = true;
+    mEtaRangeIsSet = false;
   }
 
   /// generates a single particle conforming to particle gun parameters
diff --git a/Generators/include/Generators/BoxGunParam.h b/Generators/include/Generators/BoxGunParam.h
index 716a604f5fbad..7956082646329 100644
--- a/Generators/include/Generators/BoxGunParam.h
+++ b/Generators/include/Generators/BoxGunParam.h
@@ -40,9 +40,10 @@ struct BoxGunParam : public o2::conf::ConfigurableParamHelper<BoxGunParam> {
 struct BoxGenConfig {
   int pdg = 211;                   // which particle (default pion); could make this an enum
   int number = 10;                 // how many particles
-  double eta[2] = {-1, 1};         // eta range
-  double prange[2] = {0.1, 5};     // energy range min, max in GeV
+  double eta[2] = {-1, 1};         // eta range, or rapidity range when sampleYAndPt is true
+  double prange[2] = {0.1, 5};     // p range [GeV], or pT range when sampleYAndPt is true
   double phirange[2] = {0., 360.}; // phi range
+  bool sampleYAndPt = false;       // sample uniformly in rapidity and pT instead of eta and p
 };
 
 } // end namespace eventgen
diff --git a/Generators/src/BoxGenerator.cxx b/Generators/src/BoxGenerator.cxx
index 478934d98c621..8c02addad0a00 100644
--- a/Generators/src/BoxGenerator.cxx
+++ b/Generators/src/BoxGenerator.cxx
@@ -14,6 +14,7 @@
 #include "Generators/BoxGenerator.h"
 #include "TRandom.h"
 #include "TDatabasePDG.h"
+#include <stdexcept>
 
 using namespace o2::eventgen;
 
@@ -36,7 +37,11 @@ TParticle o2::eventgen::BoxGenerator::sampleParticle() const
   // if SetCosTheta() function is used, the distribution will be uniform in
   // cos(theta)
 
-  static double mass = GetPDGMass(mPDG);
+  if (mYRangeIsSet && !mPtRangeIsSet) {
+    throw std::invalid_argument("BoxGenerator: rapidity sampling requires SetPtRange() or sampleYAndPt=true in the configuration");
+  }
+
+  const double mass = GetPDGMass(mPDG);
 
   double pabs = 0, phi, pt = 0, theta = 0, eta, y, mt, px, py, pz = 0;
   phi = gRandom->Uniform(mPhiMin, mPhiMax) * TMath::DegToRad();