feat: Log-uniform p/pt distributions for Examples (acts-project#3790)

Allows for log-uniform p/pt distributions in the Examples framework

Core/include/Acts/EventData/detail/GenerateParameters.hpp

@@ -143,6 +143,9 @@ struct GenerateQoverPOptions {
   /// Indicate if the momentum referse to transverse momentum
   bool pTransverse = true;
 
+  /// Indicate if the momentum should be uniformly distributed in log space.
+  bool pLogUniform = false;
+
   /// Charge of the parameters.
   double charge = 1;
 
@@ -157,6 +160,19 @@ inline double generateQoverP(generator_t& rng,
   using UniformIndex = std::uniform_int_distribution<std::uint8_t>;
   using UniformReal = std::uniform_real_distribution<double>;
 
+  auto drawP = [&options](generator_t& rng_, double theta_) -> double {
+    const double pTransverseScaling =
+        options.pTransverse ? 1. / std::sin(theta_) : 1.;
+
+    if (options.pLogUniform) {
+      UniformReal pLogDist(std::log(options.pMin), std::log(options.pMax));
+      return std::exp(pLogDist(rng_)) * pTransverseScaling;
+    }
+
+    UniformReal pDist(options.pMin, options.pMax);
+    return pDist(rng_) * pTransverseScaling;
+  };
+
   // choose between particle/anti-particle if requested
   // the upper limit of the distribution is inclusive
   UniformIndex particleTypeChoice(0u, options.randomizeCharge ? 1u : 0u);
@@ -165,14 +181,12 @@ inline double generateQoverP(generator_t& rng,
       options.charge,
       -options.charge,
   };
-  UniformReal pDist(options.pMin, options.pMax);
 
   // draw parameters
   const std::uint8_t type = particleTypeChoice(rng);
   const double q = qChoices[type];
 
-  const double p =
-      pDist(rng) * (options.pTransverse ? 1. / std::sin(theta) : 1.);
+  const double p = drawP(rng, theta);
   const double qOverP = (q != 0) ? q / p : 1 / p;
 
   return qOverP;

Examples/Algorithms/Generators/ActsExamples/Generators/ParametricParticleGenerator.cpp

@@ -9,57 +9,84 @@
 #include "ActsExamples/Generators/ParametricParticleGenerator.hpp"
 
 #include "Acts/Definitions/Algebra.hpp"
-#include "Acts/Definitions/Common.hpp"
 #include "Acts/Definitions/ParticleData.hpp"
 #include "Acts/Utilities/AngleHelpers.hpp"
 #include "ActsFatras/EventData/Barcode.hpp"
 #include "ActsFatras/EventData/Particle.hpp"
 
-#include <cstdint>
 #include <limits>
-#include <random>
 #include <utility>
 
 namespace ActsExamples {
 
 ParametricParticleGenerator::ParametricParticleGenerator(const Config& cfg)
     : m_cfg(cfg),
       m_charge(cfg.charge.value_or(Acts::findCharge(m_cfg.pdg).value_or(0))),
-      m_mass(cfg.mass.value_or(Acts::findMass(m_cfg.pdg).value_or(0))),
-      // since we want to draw the direction uniform on the unit sphere, we must
-      // draw from cos(theta) instead of theta. see e.g.
-      // https://mathworld.wolfram.com/SpherePointPicking.html
-      m_cosThetaMin(std::cos(m_cfg.thetaMin)),
-      // ensure upper bound is included. see e.g.
-      // https://en.cppreference.com/w/cpp/numeric/random/uniform_real_distribution
-      m_cosThetaMax(std::nextafter(std::cos(m_cfg.thetaMax),
-                                   std::numeric_limits<double>::max())),
-      // in case we force uniform eta generation
-      m_etaMin(Acts::AngleHelpers::etaFromTheta(m_cfg.thetaMin)),
-      m_etaMax(Acts::AngleHelpers::etaFromTheta(m_cfg.thetaMax)) {}
-
-std::pair<SimVertexContainer, SimParticleContainer>
-ParametricParticleGenerator::operator()(RandomEngine& rng) {
-  using UniformIndex = std::uniform_int_distribution<std::uint8_t>;
-  using UniformReal = std::uniform_real_distribution<double>;
-
-  // choose between particle/anti-particle if requested
-  // the upper limit of the distribution is inclusive
-  UniformIndex particleTypeChoice(0u, m_cfg.randomizeCharge ? 1u : 0u);
-  // (anti-)particle choice is one random draw but defines two properties
-  const Acts::PdgParticle pdgChoices[] = {
+      m_mass(cfg.mass.value_or(Acts::findMass(m_cfg.pdg).value_or(0))) {
+  m_pdgChoices = {
       m_cfg.pdg,
       static_cast<Acts::PdgParticle>(-m_cfg.pdg),
   };
-  const double qChoices[] = {
+  m_qChoices = {
       m_charge,
       -m_charge,
   };
-  UniformReal phiDist(m_cfg.phiMin, m_cfg.phiMax);
-  UniformReal cosThetaDist(m_cosThetaMin, m_cosThetaMax);
-  UniformReal etaDist(m_etaMin, m_etaMax);
-  UniformReal pDist(m_cfg.pMin, m_cfg.pMax);
 
+  // choose between particle/anti-particle if requested
+  // the upper limit of the distribution is inclusive
+  m_particleTypeChoice = UniformIndex(0u, m_cfg.randomizeCharge ? 1u : 0u);
+  m_phiDist = UniformReal(m_cfg.phiMin, m_cfg.phiMax);
+
+  if (m_cfg.etaUniform) {
+    double etaMin = Acts::AngleHelpers::etaFromTheta(m_cfg.thetaMin);
+    double etaMax = Acts::AngleHelpers::etaFromTheta(m_cfg.thetaMax);
+
+    UniformReal etaDist(etaMin, etaMax);
+
+    m_sinCosThetaDist =
+        [=](RandomEngine& rng) mutable -> std::pair<double, double> {
+      const double eta = etaDist(rng);
+      const double theta = Acts::AngleHelpers::thetaFromEta(eta);
+      return {std::sin(theta), std::cos(theta)};
+    };
+  } else {
+    // since we want to draw the direction uniform on the unit sphere, we must
+    // draw from cos(theta) instead of theta. see e.g.
+    // https://mathworld.wolfram.com/SpherePointPicking.html
+    double cosThetaMin = std::cos(m_cfg.thetaMin);
+    // ensure upper bound is included. see e.g.
+    // https://en.cppreference.com/w/cpp/numeric/random/uniform_real_distribution
+    double cosThetaMax = std::nextafter(std::cos(m_cfg.thetaMax),
+                                        std::numeric_limits<double>::max());
+
+    UniformReal cosThetaDist(cosThetaMin, cosThetaMax);
+
+    m_sinCosThetaDist =
+        [=](RandomEngine& rng) mutable -> std::pair<double, double> {
+      const double cosTheta = cosThetaDist(rng);
+      return {std::sqrt(1 - cosTheta * cosTheta), cosTheta};
+    };
+  }
+
+  if (m_cfg.pLogUniform) {
+    // distributes p or pt uniformly in log space
+    UniformReal dist(std::log(m_cfg.pMin), std::log(m_cfg.pMax));
+
+    m_somePDist = [=](RandomEngine& rng) mutable -> double {
+      return std::exp(dist(rng));
+    };
+  } else {
+    // distributes p or pt uniformly
+    UniformReal dist(m_cfg.pMin, m_cfg.pMax);
+
+    m_somePDist = [=](RandomEngine& rng) mutable -> double {
+      return dist(rng);
+    };
+  }
+}
+
+std::pair<SimVertexContainer, SimParticleContainer>
+ParametricParticleGenerator::operator()(RandomEngine& rng) {
   SimVertexContainer::sequence_type vertices;
   SimParticleContainer::sequence_type particles;
 
@@ -74,33 +101,22 @@ ParametricParticleGenerator::operator()(RandomEngine& rng) {
     primaryVertex.outgoing.insert(pid);
 
     // draw parameters
-    const unsigned int type = particleTypeChoice(rng);
-    const Acts::PdgParticle pdg = pdgChoices[type];
-    const double q = qChoices[type];
-    const double phi = phiDist(rng);
-    double p = pDist(rng);
-
-    // we already have sin/cos theta. they can be used directly to
-    Acts::Vector3 dir;
-    double cosTheta = 0.;
-    double sinTheta = 0.;
-    if (!m_cfg.etaUniform) {
-      cosTheta = cosThetaDist(rng);
-      sinTheta = std::sqrt(1 - cosTheta * cosTheta);
-    } else {
-      const double eta = etaDist(rng);
-      const double theta = Acts::AngleHelpers::thetaFromEta(eta);
-      sinTheta = std::sin(theta);
-      cosTheta = std::cos(theta);
-    }
-    dir[Acts::eMom0] = sinTheta * std::cos(phi);
-    dir[Acts::eMom1] = sinTheta * std::sin(phi);
-    dir[Acts::eMom2] = cosTheta;
+    const unsigned int type = m_particleTypeChoice(rng);
+    const Acts::PdgParticle pdg = m_pdgChoices[type];
+    const double q = m_qChoices[type];
+    const double phi = m_phiDist(rng);
+    const double someP = m_somePDist(rng);
+
+    const auto [sinTheta, cosTheta] = m_sinCosThetaDist(rng);
+    // we already have sin/cos theta. they can be used directly
+    const Acts::Vector3 dir = {sinTheta * std::cos(phi),
+                               sinTheta * std::sin(phi), cosTheta};
+
+    const double p = someP * (m_cfg.pTransverse ? 1. / sinTheta : 1.);
 
     // construct the particle;
     ActsFatras::Particle particle(pid, pdg, q, m_mass);
     particle.setDirection(dir);
-    p *= m_cfg.pTransverse ? 1. / sinTheta : 1.;
     particle.setAbsoluteMomentum(p);
 
     // generated particle ids are already ordered and should end up at the end

Examples/Algorithms/Generators/ActsExamples/Generators/ParametricParticleGenerator.hpp

@@ -15,10 +15,10 @@
 #include "ActsExamples/Generators/EventGenerator.hpp"
 
 #include <array>
-#include <cmath>
 #include <cstddef>
 #include <limits>
 #include <optional>
+#include <random>
 
 namespace ActsExamples {
 
@@ -49,8 +49,10 @@ class ParametricParticleGenerator : public EventGenerator::ParticlesGenerator {
     /// Low, high (exclusive) for absolute/transverse momentum.
     double pMin = 1 * Acts::UnitConstants::GeV;
     double pMax = 10 * Acts::UnitConstants::GeV;
-    /// Indicate if the momentum referse to transverse momentum
+    /// Indicate if the momentum refers to transverse momentum.
     bool pTransverse = false;
+    /// Indicate if the momentum should be uniformly distributed in log space.
+    bool pLogUniform = false;
     /// (Absolute) PDG particle number to identify the particle type.
     Acts::PdgParticle pdg = Acts::PdgParticle::eMuon;
     /// Randomize the charge and flip the PDG particle number sign accordingly.
@@ -71,14 +73,23 @@ class ParametricParticleGenerator : public EventGenerator::ParticlesGenerator {
       RandomEngine& rng) override;
 
  private:
+  using UniformIndex = std::uniform_int_distribution<std::uint8_t>;
+  using UniformReal = std::uniform_real_distribution<double>;
+
   Config m_cfg;
+
   // will be automatically set from PDG data tables
-  double m_charge;
-  double m_mass;
-  double m_cosThetaMin;
-  double m_cosThetaMax;
-  double m_etaMin;
-  double m_etaMax;
+  double m_charge{};
+  double m_mass{};
+
+  // (anti-)particle choice is one random draw but defines two properties
+  std::array<Acts::PdgParticle, 2> m_pdgChoices{};
+  std::array<double, 2> m_qChoices{};
+
+  UniformIndex m_particleTypeChoice;
+  UniformReal m_phiDist;
+  std::function<std::pair<double, double>(RandomEngine& rng)> m_sinCosThetaDist;
+  std::function<double(RandomEngine& rng)> m_somePDist;
 };
 
 }  // namespace ActsExamples

Examples/Python/python/acts/examples/simulation.py

@@ -18,8 +18,8 @@
 # Examples/Algorithms/Generators/ActsExamples/Generators/ParametricParticleGenerator.hpp
 MomentumConfig = namedtuple(
     "MomentumConfig",
-    ["min", "max", "transverse"],
-    defaults=[None, None, None],
+    ["min", "max", "transverse", "logUniform"],
+    defaults=[None, None, None, None],
 )
 EtaConfig = namedtuple(
     "EtaConfig", ["min", "max", "uniform"], defaults=[None, None, None]
@@ -80,8 +80,8 @@ def addParticleGun(
         the output folder for the Csv output, None triggers no output
     outputDirRoot : Path|str, path, None
         the output folder for the Root output, None triggers no output
-    momentumConfig : MomentumConfig(min, max, transverse)
-        momentum configuration: minimum momentum, maximum momentum, transverse
+    momentumConfig : MomentumConfig(min, max, transverse, logUniform)
+        momentum configuration: minimum momentum, maximum momentum, transverse, log-uniform
     etaConfig : EtaConfig(min, max, uniform)
         pseudorapidity configuration: eta min, eta max, uniform
     phiConfig : PhiConfig(min, max)
@@ -118,6 +118,7 @@ def addParticleGun(
                     **acts.examples.defaultKWArgs(
                         p=(momentumConfig.min, momentumConfig.max),
                         pTransverse=momentumConfig.transverse,
+                        pLogUniform=momentumConfig.logUniform,
                         eta=(etaConfig.min, etaConfig.max),
                         phi=(phiConfig.min, phiConfig.max),
                         etaUniform=etaConfig.uniform,

Examples/Python/src/Generators.cpp

@@ -176,6 +176,7 @@ void addGenerators(Context& ctx) {
         .def_readwrite("pMin", &Config::pMin)
         .def_readwrite("pMax", &Config::pMax)
         .def_readwrite("pTransverse", &Config::pTransverse)
+        .def_readwrite("pLogUniform", &Config::pLogUniform)
         .def_readwrite("pdg", &Config::pdg)
         .def_readwrite("randomizeCharge", &Config::randomizeCharge)
         .def_readwrite("numParticles", &Config::numParticles)