generators/flowAfterburner/flowAfterburner.cc

0001 // File:  Generators/FlowAfterburnber/AddFlowByShifting.cxx
0002 // Description:
0003 //    This code is used to introduce particle flow
0004 //    to particles from generated events
0005 //
0006 // AuthorList:
0007 // Andrzej Olszewski: Initial Code February 2006
0008 // 11.10.2006: Add predefined flow function by name
0009
0010 // The initialization of this class is trivial.  There's really no
0011 // need for it.  Pass in a pointer to a random generator, algorithm
0012 // selection and an event.
0013
0014 #include "flowAfterburner.h"
0015 #include "AfterburnerAlgo.h"
0016
0017 #include <phool/phool.h>
0018
0019 #include <gsl/gsl_errno.h>
0020 #include <gsl/gsl_math.h>
0021 #include <gsl/gsl_roots.h>
0022
0023 #include <HepMC/GenEvent.h>
0024 #include <HepMC/GenParticle.h>  // for GenParticle
0025 #include <HepMC/GenRanges.h>
0026 #include <HepMC/GenVertex.h>      // for GenVertex, GenVertex::part...
0027 #include <HepMC/HeavyIon.h>       // for HeavyIon
0028 #include <HepMC/IteratorRange.h>  // for children, descendants
0029 #include <HepMC/SimpleVector.h>   // for FourVector
0030
0031 #include <CLHEP/Random/RandFlat.h>
0032 #include <CLHEP/Vector/LorentzVector.h>
0033 #include <CLHEP/Random/MTwistEngine.h>
0034
0035 #include <cmath>
0036 #include <cstdlib>  // for exit
0037 #include <iostream>
0038 #include <map>  // for map
0039 #include <ctime>  // for time
0040 #include <string>
0041 #include <algorithm>  // for max, min
0042
0043 namespace CLHEP
0044 {
0045   class HepRandomEngine;
0046 }
0047
0048 Afterburner::Afterburner( const std::string &algorithmName,
0049                           CLHEP::HepRandomEngine *engine,
0050                           float mineta, float maxeta,
0051                           float minpt, float maxpt )
0052   : m_algo(new AfterburnerAlgo(AfterburnerAlgo::getAlgoFromName(algorithmName)))
0053   , m_engine(engine)
0054   , m_ownAlgo(true)
0055   , m_ownEngine(engine == nullptr)
0056   , m_mineta(mineta)
0057   , m_maxeta(maxeta)
0058   , m_minpt(minpt)
0059   , m_maxpt(maxpt)
0060   , m_phishift(0.0)
0061 {
0062
0063     if (!m_engine)
0064     {
0065       // seed with time if no engine is provided
0066       m_engine = new CLHEP::MTwistEngine(static_cast<long>(time(nullptr)));
0067     }
0068
0069 }
0070
0071
0072 Afterburner::~Afterburner()
0073 {
0074   if (m_ownAlgo)
0075   {
0076     delete m_algo;
0077   }
0078   if (m_ownEngine)
0079   {
0080     delete m_engine;
0081   }
0082 }
0083
0084 Afterburner::Afterburner(Afterburner&& o) noexcept
0085   : m_algo(o.m_algo)
0086   , m_engine(o.m_engine)
0087   , m_ownAlgo(o.m_ownAlgo)
0088   , m_ownEngine(o.m_ownEngine)
0089   , m_mineta(o.m_mineta)
0090   , m_maxeta(o.m_maxeta)
0091   , m_minpt(o.m_minpt)
0092   , m_maxpt(o.m_maxpt)
0093   , m_phishift(o.m_phishift)
0094 {
0095   std::copy(std::begin(o.m_psi_n), std::end(o.m_psi_n), std::begin(m_psi_n));
0096   o.m_algo = nullptr;
0097   o.m_engine = nullptr;
0098   o.m_ownAlgo = false;
0099   o.m_ownEngine = false;
0100 }
0101
0102
0103 Afterburner& Afterburner::operator=(Afterburner&& o) noexcept
0104 {
0105   if (this != &o) {
0106     if (m_ownAlgo)
0107     {
0108       delete m_algo;
0109     }
0110     if (m_ownEngine)
0111     {
0112       delete m_engine;
0113     }
0114     m_algo = o.m_algo;
0115     o.m_algo = nullptr;
0116     m_engine = o.m_engine;
0117     o.m_engine = nullptr;
0118     m_ownAlgo = o.m_ownAlgo;
0119     o.m_ownAlgo = false;
0120     m_ownEngine = o.m_ownEngine;
0121     o.m_ownEngine = false;
0122     m_mineta = o.m_mineta;
0123     m_maxeta = o.m_maxeta;
0124     m_minpt  = o.m_minpt;
0125     m_maxpt  = o.m_maxpt;
0126     m_phishift = o.m_phishift;
0127     std::copy(std::begin(o.m_psi_n), std::end(o.m_psi_n), std::begin(m_psi_n));
0128   }
0129   return *this;
0130 }
0131
0132
0133 void Afterburner::setAlgo(AfterburnerAlgo::flowAfterburnerAlgorithm algo_type)
0134 {
0135   if (m_ownAlgo && m_algo)
0136   {
0137     delete m_algo;
0138   }
0139   m_algo = new AfterburnerAlgo(algo_type);
0140   m_ownAlgo = true;
0141 }
0142
0143 void Afterburner::setAlgo(const std::string &name)
0144 {
0145   setAlgo(AfterburnerAlgo::getAlgoFromName(name));
0146 }
0147 void Afterburner::setAlgo(AfterburnerAlgo* algo)
0148 {
0149   if (m_ownAlgo && m_algo)
0150   {
0151     delete m_algo;
0152   }
0153   m_algo = algo;
0154   m_ownAlgo = false; // external ownership
0155 }
0156
0157 void Afterburner::setEngine(CLHEP::HepRandomEngine* engine)
0158 {
0159   if (m_ownEngine && m_engine)
0160   {
0161     delete m_engine;
0162   }
0163   m_engine = engine;
0164   m_ownEngine = false; // external ownership
0165 }
0166
0167 void Afterburner::setEtaRange(float mineta, float maxeta)
0168 {
0169   m_mineta = std::min(mineta, maxeta);
0170   m_maxeta = std::max(mineta, maxeta);
0171 }
0172
0173 void Afterburner::setPtRange(float minpt, float maxpt)
0174 {
0175   m_minpt = std::min(minpt, maxpt);
0176   m_maxpt = std::max(minpt, maxpt);
0177 }
0178
0179 void Afterburner::setPsiN(unsigned int n, float psi)
0180 {
0181   if (n < 1 || n > 6)
0182   {
0183     std::cout << PHWHERE << ": Flow Afterburner set reaction plane angle psi_n for n=" << n << " which is out of range. Ignoring." << std::endl;
0184     return;
0185   }
0186   m_psi_n[n - 1] = psi;
0187 }
0188
0189 float Afterburner::getPsiN(unsigned int n) const
0190 {
0191   if (n < 1 || n > 6)
0192   {
0193     std::cout << PHWHERE << ": Flow Afterburner requested reaction plane angle psi_n for n=" << n << " which is out of range. Returning 0.0" << std::endl;
0194     return 0.0;
0195   }
0196   return m_psi_n[n - 1];
0197
0198 }
0199
0200 double Afterburner::vn_func(double x, void* params)
0201 {
0202   const float* par = static_cast<const float*>(params);
0203   const double phi0 = par[0];
0204   const float* vn   = par + 1;
0205   const float* psi  = par + 7;
0206   double s = 0.0;
0207   for (int n = 1; n <= 6; ++n) {
0208     s += vn[n-1] * std::sin(static_cast<double>(n) * (x - psi[n-1])) / static_cast<double>(n);
0209   }
0210   return (x + 2.0 * s) - phi0;
0211 }
0212
0213
0214 void Afterburner::throw_psi_n(HepMC::GenEvent *event)
0215 {
0216   HepMC::HeavyIon *hi = event->heavy_ion();
0217   if (!hi)
0218   {
0219     std::cout << PHWHERE << ": Flow Afterburner needs the Heavy Ion Event Info, GenEvent::heavy_ion() returns NULL" << std::endl;
0220     exit(1);
0221   }
0222   float psi_n[6] = {0.0, 0.0, 0.0, 0.0, 0.0, 0.0}; // reaction plane angles
0223   for (int i = 0; i < 6; i++)
0224   {
0225     // Principal value must be within -PI/n to PI/n
0226     psi_n[i] = ((CLHEP::RandFlat::shoot(m_engine) - 0.5) * 2 * M_PI )
0227                 /( static_cast<float>(i + 1) );
0228   }
0229
0230   // The psi2 plane is aligned with the impact parameter
0231   psi_n[1] = hi->event_plane_angle();
0232
0233   // Ensure that Psi2 is within [-PI/2,PI/2]
0234   psi_n[1] = std::atan2(std::sin(2.0 * psi_n[1]), std::cos(2.0 * psi_n[1])) / 2.0;
0235
0236
0237   for (unsigned int i = 1; i <= 6; ++i)
0238   {
0239     setPsiN(i, psi_n[i - 1]);
0240   }
0241   return;
0242 }
0243
0244 void Afterburner::AddFlowToParentAndMoveDescendants(HepMC::GenEvent *event, HepMC::GenParticle *parent )
0245 {
0246   CLHEP::HepLorentzVector momentum(parent->momentum().px(),
0247                                     parent->momentum().py(),
0248                                     parent->momentum().pz(),
0249                                     parent->momentum().e());
0250   double pt = momentum.perp();
0251   double eta = momentum.pseudoRapidity();
0252   double phi_0 = momentum.phi();
0253
0254   HepMC::HeavyIon* hi = event->heavy_ion();
0255   if (!hi)
0256   {
0257     std::cout << PHWHERE << ": HeavyIon info missing in GenEvent. Cannot apply flow." << std::endl;
0258     std::exit(1);
0259   }
0260
0261   m_algo->set_impact_parameter(hi->impact_parameter());
0262   m_algo->calc_flow(eta, pt, m_engine); // add engine for fluctuations (if enabled)
0263
0264   const gsl_root_fsolver_type *T = gsl_root_fsolver_brent;
0265   gsl_root_fsolver *s = gsl_root_fsolver_alloc(T);
0266   double x_lo = -2 * M_PI;
0267   double x_hi = 2 * M_PI;
0268   float params[13] = {};
0269   params[0]  = static_cast<float>(phi_0);
0270   for (int i = 0; i < 6; ++i)
0271   {
0272     params[1 + i] = m_algo->get_vn(i+1);
0273     params[7 + i] = getPsiN(i+1);
0274   }
0275
0276   gsl_function F;
0277   F.function = &Afterburner::vn_func;
0278   F.params = params;
0279   gsl_root_fsolver_set(s, &F, x_lo, x_hi);
0280
0281   int status;
0282   int iter = 0;
0283   double phi = 0;
0284   do
0285   {
0286     ++iter;
0287     gsl_root_fsolver_iterate(s);
0288     phi  = gsl_root_fsolver_root(s);
0289     x_lo = gsl_root_fsolver_x_lower(s);
0290     x_hi = gsl_root_fsolver_x_upper(s);
0291     status = gsl_root_test_interval(x_lo, x_hi, 0, 1e-5);
0292   } while (status == GSL_CONTINUE && iter < 1000);
0293
0294   if (iter >= 1000) {
0295     gsl_root_fsolver_free(s);
0296     m_phishift = 0.0;
0297     return; // do not rotate anything on failure
0298   }
0299
0300   gsl_root_fsolver_free(s);
0301
0302   m_phishift = phi - phi_0;
0303   if (fabs(m_phishift) > 1e-7)
0304   {
0305     momentum.rotateZ(m_phishift);  // DPM check units * Gaudi::Units::rad);
0306     parent->set_momentum(momentum);
0307   }
0308
0309   HepMC::GenVertex *endvtx = parent->end_vertex();
0310   if (endvtx)
0311   {
0312     // now rotate descendant vertices
0313     for (HepMC::GenVertex::vertex_iterator descvtxit = endvtx->vertices_begin(HepMC::descendants);
0314           descvtxit != endvtx->vertices_end(HepMC::descendants);
0315           ++descvtxit)
0316     {
0317       HepMC::GenVertex *descvtx = (*descvtxit);
0318
0319       // rotate vertex (magic number?)
0320       if (fabs(m_phishift) > 1e-7)
0321       {
0322         CLHEP::HepLorentzVector position(descvtx->position().x(),
0323                                           descvtx->position().y(),
0324                                           descvtx->position().z(),
0325                                           descvtx->position().t());
0326         position.rotateZ(m_phishift);  // DPM check units
0327         descvtx->set_position(position);
0328       }
0329
0330       // now rotate their associated particles
0331       for (HepMC::GenVertex::particle_iterator descpartit = descvtx->particles_begin(HepMC::children);
0332             descpartit != descvtx->particles_end(HepMC::children);
0333             ++descpartit)
0334       {
0335         HepMC::GenParticle *descpart = (*descpartit);
0336         CLHEP::HepLorentzVector desmomentum(descpart->momentum().px(),
0337                                           descpart->momentum().py(),
0338                                           descpart->momentum().pz(),
0339                                           descpart->momentum().e());
0340         // Rotate particle
0341         if (fabs(m_phishift) > 1e-7)
0342         {
0343           desmomentum.rotateZ(m_phishift);  // DPM check units * Gaudi::Units::rad);
0344           descpart->set_momentum(desmomentum);
0345         }
0346       } // end of particles loop
0347     } // end of vertices loop
0348   } // end of if (endvtx)
0349
0350
0351   return ;
0352 }
0353
0354 void Afterburner::readLegacyArguments(
0355     CLHEP::HepRandomEngine *engine,
0356     const std::string &algorithmName,
0357     float mineta, float maxeta,
0358     float minpt, float maxpt)
0359 {
0360   if (engine)
0361   {
0362     setEngine(engine);
0363   }
0364   if (!algorithmName.empty())
0365   {
0366     setAlgo(algorithmName);
0367   }
0368   if (mineta != -5.0 || maxeta != 5.0)
0369   {
0370     setEtaRange(mineta, maxeta);
0371   }
0372   if (minpt != 0.0 || maxpt != 100.0)
0373   {
0374     setPtRange(minpt, maxpt);
0375   }
0376
0377   return;
0378 }
0379
0380 int Afterburner::flowAfterburner(HepMC::GenEvent *event,
0381                                   CLHEP::HepRandomEngine *engine,
0382                                   const std::string &algorithmName,
0383                                   float mineta, float maxeta,
0384                                   float minpt, float maxpt )
0385 {
0386   // check legacy arguments to see if they are set
0387   readLegacyArguments(engine, algorithmName, mineta, maxeta, minpt, maxpt);
0388
0389   throw_psi_n(event); // throw reaction plane angles
0390
0391   HepMC::GenVertex *mainvtx = event->barcode_to_vertex(-1);
0392   if (!mainvtx)
0393   {
0394     std::cout << PHWHERE << ": Flow Afterburner cannot find main vertex in GenEvent. Cannot apply flow." << std::endl;
0395     return -1;
0396   }
0397
0398   // Loop over all children of this vertex
0399   HepMC::GenVertexParticleRange r(*mainvtx, HepMC::children);
0400
0401   for (HepMC::GenVertex::particle_iterator it = r.begin(); it != r.end(); it++)
0402   {
0403     // Process particles from main vertex
0404     HepMC::GenParticle *parent = (*it);
0405
0406     CLHEP::HepLorentzVector momentum(parent->momentum().px(),
0407                                      parent->momentum().py(),
0408                                      parent->momentum().pz(),
0409                                      parent->momentum().e());
0410
0411     float eta = momentum.pseudoRapidity();
0412     if (eta < m_mineta || eta > m_maxeta)
0413     {
0414       continue;
0415     }
0416
0417     float pT = momentum.perp();
0418     if (pT < m_minpt || pT > m_maxpt)
0419     {
0420       continue;
0421     }
0422
0423     AddFlowToParentAndMoveDescendants(event, parent);
0424   } // end of particles loop
0425
0426   return 0;
0427
0428 }
0429
0430 // Legacy function, use the Afterburner class instead
0431 int flowAfterburner(HepMC::GenEvent *event,
0432                     CLHEP::HepRandomEngine *engine,
0433                     const std::string &algorithmName,
0434                     float mineta, float maxeta,
0435                     float minpt, float maxpt )
0436 {
0437   Afterburner afterburner(algorithmName, engine, mineta, maxeta, minpt, maxpt);
0438   return afterburner.flowAfterburner(event, engine, algorithmName, mineta, maxeta, minpt, maxpt);
0439 }
0440
0441