insert_muon/analysis/gen_particles.cxx

0001 #include "HepMC3/GenEvent.h"
0002 #include "HepMC3/ReaderAscii.h"
0003 #include "HepMC3/WriterAscii.h"
0004 #include "HepMC3/Print.h"
0005
0006 #include "TRandom3.h"
0007 #include "TVector3.h"
0008
0009 #include <iostream>
0010 #include <random>
0011 #include <cmath>
0012 #include <math.h>
0013 #include <TMath.h>
0014 #include <TDatabasePDG.h>
0015 #include <TParticlePDG.h>
0016
0017 using namespace HepMC3;
0018
0019 // Generate single electron as input to the Insert simulation.
0020 // --
0021 // We generate events with a constant polar angle with respect to
0022 // the proton direction and then rotate so that the events are given
0023 // in normal lab coordinate system
0024 // --
0025 void gen_particles(
0026                     int n_events = 1000,
0027                     const char* out_fname = "gen_particles.hepmc",
0028                     TString particle_name = "e-",
0029                     double th_min = 3., // Minimum polar angle, in degrees
0030             double th_max = 3., // Maximum polar angle, in degrees
0031             double phi_min = 0., // Minimum azimuthal angle, in degrees
0032                     double phi_max = 360., // Maximum azimuthal angle, in degrees
0033                     double p = 10.,  // Momentum in GeV/c
0034             int dist = 0,  //Momentum distribution: 0=fixed, 1=uniform, 2=Gaussian
0035             int useCrossingAngle=1  // 0= no rotation, 1 = -25 mrad
0036                   )
0037 {
0038   WriterAscii hepmc_output(out_fname);
0039   int events_parsed = 0;
0040   GenEvent evt(Units::GEV, Units::MM);
0041
0042   // Random number generator
0043   TRandom3 *r1 = new TRandom3(0); //Use time as random seed
0044
0045   // Getting generated particle information
0046   TDatabasePDG *pdg = new TDatabasePDG();
0047   TParticlePDG *particle = pdg->GetParticle(particle_name);
0048   const double mass = particle->Mass();
0049   const int pdgID = particle->PdgCode();
0050
0051   for (events_parsed = 0; events_parsed < n_events; events_parsed++) {
0052
0053     //Set the event number
0054     evt.set_event_number(events_parsed);
0055
0056     // FourVector(px,py,pz,e,pdgid,status)
0057     // type 4 is beam
0058     // pdgid 11 - electron
0059     // pdgid 111 - pi0
0060     // pdgid 2212 - proton
0061     GenParticlePtr p1 =
0062         std::make_shared<GenParticle>(FourVector(0.0, 0.0, 10.0, 10.0), 11, 4);
0063     GenParticlePtr p2 = std::make_shared<GenParticle>(
0064         FourVector(0.0, 0.0, 0.0, 0.938), 2212, 4);
0065
0066     // Define momentum with respect to proton direction
0067     double phi   = r1->Uniform(phi_min*TMath::DegToRad(),phi_max*TMath::DegToRad());
0068     double th    = r1->Uniform(th_min*TMath::DegToRad(),th_max*TMath::DegToRad());
0069
0070     //Total momentum distribution
0071     double pevent = -1;
0072     if(dist==0){ //fixed
0073     pevent = p;
0074     }
0075     else if(dist==1){ //Uniform: +-50% variation
0076     pevent = p*(1. + r1->Uniform(-0.5,0.5) );
0077     }
0078     else if(dist==2){  //Gaussian: Sigma = 0.1*mean
0079     while(pevent<0) //Avoid negative values
0080         pevent = r1->Gaus(p,0.1*p);
0081     }
0082
0083     double px    = pevent * std::cos(phi) * std::sin(th);
0084     double py    = pevent * std::sin(phi) * std::sin(th);
0085     double pz    = pevent * std::cos(th);
0086     TVector3 pvec(px,py,pz);
0087
0088     //Rotate to lab coordinate system
0089     double cross_angle = -25./1000.*useCrossingAngle; //in Rad
0090     TVector3 pbeam_dir(sin(cross_angle),0,cos(cross_angle)); //proton beam direction
0091     pvec.RotateY(-pbeam_dir.Theta()); // Theta is returned positive, beam in negative X
0092     // type 1 is final state
0093     // pdgid 11 - electron 0.510 MeV/c^2
0094     GenParticlePtr p3 = std::make_shared<GenParticle>(
0095         FourVector(
0096             pvec.X(), pvec.Y(), pvec.Z(),
0097             sqrt(pevent*pevent + (mass * mass))),
0098         pdgID, 1);
0099
0100     //If wanted, set non-zero vertex
0101     double vx = 0.;
0102     double vy = 0.;
0103     double vz = 0.;
0104     double vt = 0.;
0105
0106     GenVertexPtr v1 = std::make_shared<GenVertex>();
0107     evt.shift_position_by(FourVector(vx, vy, vz, vt));
0108     v1->add_particle_in(p1);
0109     v1->add_particle_in(p2);
0110
0111     v1->add_particle_out(p3);
0112     evt.add_vertex(v1);
0113
0114     if (events_parsed == 0) {
0115       std::cout << "First event: " << std::endl;
0116       Print::listing(evt);
0117     }
0118
0119     hepmc_output.write_event(evt);
0120     if (events_parsed % 100 == 0) {
0121       std::cout << "Event: " << events_parsed << std::endl;
0122     }
0123     evt.clear();
0124   }
0125   hepmc_output.close();
0126   std::cout << "Events parsed and written: " << events_parsed << std::endl;
0127 }