BoxGenerator: enable sampling of pT and rapidity instead of p and eta

Author: fmazzasc

Generators/include/Generators/BoxGenerator.h

@@ -35,32 +35,47 @@ 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,
                double etamax,
                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)
   {
     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

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

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,6 +37,10 @@ TParticle o2::eventgen::BoxGenerator::sampleParticle() const
   // if SetCosTheta() function is used, the distribution will be uniform in
   // cos(theta)
 
+  if (mYRangeIsSet && !mPtRangeIsSet) {
+    throw std::invalid_argument("BoxGenerator: rapidity sampling requires SetPtRange() or sampleYAndPt=true in the configuration");
+  }
+
   // per instance, since several box generators with different PDG codes can coexist
   const double mass = GetPDGMass(mPDG);