Pythia8Generator.cpp

#include "pxl/hep.hh"
#include "pxl/core.hh"
#include "pxl/core/macros.hh"
#include "pxl/core/PluginManager.hh"
#include "pxl/modules/Module.hh"
#include "pxl/modules/ModuleFactory.hh"
#include "pxl/modules.hh"
#include "pxl/core/Event.hh"
#include "pxl/hep/EventView.hh"

#include "Pythia.h"

#include<cmath>
#include<string>
#include<vector>
#include <map>
#include <algorithm>

#include <sstream>

static pxl::Logger logger("Pythia8Generator");

class Pythia8Generator : public pxl::Module
{
private:
    Pythia8::Pythia pythia8;
    int64_t counter;
    std::vector<std::string> coll_particle1;
    std::vector<std::string> coll_particle2;
    double cme;
    int64_t seed;
    int64_t numberOfObjects;
    bool running;

    std::vector<std::string> mode;
    std::vector<std::string> colliding_particles;


    enum StoredParticles {
    	ALL,
    	ALL_STABLE,
    	HARD_PROCESS
    };
    StoredParticles storedParticles;

public:
    
        // every Module needs a unique type
        static const std::string &getStaticType()
        {
            static std::string type ("Pythia8Generator");
            return type;
        }
        
        // static and dynamic methods are needed
        const std::string &getType() const
        {
            return getStaticType();
        }

        
        Pythia8Generator(): Module()
        {
        	colliding_particles.push_back("proton");
        	colliding_particles.push_back("anti proton");
        	colliding_particles.push_back("electron");
        	colliding_particles.push_back("anti electron");
        	colliding_particles.push_back("muon");
        	colliding_particles.push_back("anti muon");
        	colliding_particles.push_back("proton");

        	mode.push_back("only hard process");
        	mode.push_back("all stable");
        	mode.push_back("all");
        	mode.push_back("all");
            addSource("out","");
            addOption("seed","",(int64_t)123);
            addOption("number"," ",(int64_t)10);
            addOption("colliding particle 1","the first particle for collision",colliding_particles,pxl::OptionDescription::USAGE_TEXT_SELECTION);
            addOption("colliding particle 2","the second particle for collision",colliding_particles,pxl::OptionDescription::USAGE_TEXT_SELECTION);
            addOption("center mass energy","",200.0);
            addOption("eventview name","the name of the eventview","Generated");
            addOption("stored particles","the particles which are stored into a new eventview with the given name",mode,pxl::OptionDescription::USAGE_TEXT_SELECTION);
            running = true;
        }
        
		int64_t getParticleId(std::string name)
		{
			if (name=="proton") {
				return 2212;
			}
			if (name=="anti proton") {
				return -2212;
			}
			if (name=="electron") {
				return 11;
			}
			if (name=="anti electron") {
				return -11;
			}
			if (name=="muon") {
				return 13;
			}
			if (name=="anti muon") {
				return -13;
			}
		}

		StoredParticles getStoredParticles(std::string name)
		{
			if (name=="only hard process"){
				return HARD_PROCESS;
			}
			if (name=="all stable"){
				return ALL_STABLE;
			}
			if (name=="all"){
				return ALL;
			}
		}

        ~Pythia8Generator()
        {
        }

        void beginJob() 
        {
            getOption("number",numberOfObjects);
			getOption("seed",seed);
            pythia8.readString("Random:setSeed = on");
            pythia8.readString("Random:seed = "+seed);
            pythia8.readString("WeakZ0:gmZmode=2");
            pythia8.readString("WeakBosonExchange:all=on");
            pythia8.readString("WeakSingleBoson:all=on");
            pythia8.readString("WeakBosonAndParton:all=on");
            pythia8.readString("PhaseSpace:mHatMin=50");
            pythia8.readString("PhaseSpace:mHatMax=150");
            getOption("colliding particle 1",coll_particle1);
            getOption("colliding particle 2",coll_particle2);
            getOption("center mass energy",cme);
            pythia8.init(getParticleId(coll_particle1.back()), getParticleId(coll_particle2.back()), cme);
            getOption("stored particles",mode);
            storedParticles=getStoredParticles(mode.back());
        }
        
        void beginRun()
        {
        }
        
        bool isRunnable() const
        {
            return running;
        }
        
        bool analyse(pxl::Sink* sink)
        {
            if (counter < numberOfObjects)
            {
                counter+=1;
                pythia8.next();
                Pythia8::Event pEvent = pythia8.event;
                int num = pEvent.size();
                pxl::Event* event = new pxl::Event();
                pxl::EventView* eventView = new pxl::EventView();
                eventView->setName("Generated");
                event->insertObject(eventView);
                map<int, pxl::Particle*> idMap;
                for (int cnt=0; cnt<num;++cnt) {
                    Pythia8::Particle particle = pEvent[cnt];
                    int status = particle.status();
                    bool accepted = false;
                    if (storedParticles==HARD_PROCESS) {
                    	accepted=(fabs(status)>=21 and fabs(status)<=29);
                    }
                    if (storedParticles==ALL_STABLE) {
						accepted=particle.isFinal();
					}
                    if (storedParticles==ALL) {
						accepted=true;
					}
                    if (accepted)
                    {
                        pxl::Particle* pxl_part = new pxl::Particle();
                        pxl_part->setName(particle.name());
                        pxl_part->setP4(particle.px(),particle.py(),particle.pz(),particle.e());
                        pxl_part->setUserRecord("status",status);
                        pxl_part->setUserRecord("scale",particle.scale());
                        pxl::Basic3Vector production_vertex(particle.xProd(),particle.yProd(),particle.zProd());
                        pxl_part->setUserRecord("production_vertex",production_vertex);
                        pxl_part->setCharge(particle.charge());
                        eventView->insertObject(pxl_part);
                        idMap[cnt]=pxl_part;
                    }
                }
                map<int, pxl::Particle*>::iterator it;
                for (int cnt=0; cnt<num;++cnt) {
                    Pythia8::Particle particle = pEvent[cnt];
                    if (idMap.find(cnt)!=idMap.end()) {
                        
                        int mother1ID = particle.mother1();
                        if (mother1ID>0) {
                            if (idMap.find(mother1ID)!=idMap.end()) {
                                (idMap[mother1ID])->linkDaughter(idMap[cnt]);
                            }
                        }
                        int mother2ID = particle.mother2();
                        if (mother2ID>0) {
                            if (idMap.find(mother2ID)!=idMap.end()) {
                                (idMap[mother2ID])->linkDaughter(idMap[cnt]);
                            }
                        }
                    }
                }
                
                (getSource("out"))->setTargets(event);
                return (getSource("out"))->processTargets();
            }
            running = false;
            return false;
        }
        
        void endRun()
        {
        }

        void endJob()
        {
        }
        
};

PXL_MODULE_INIT(Pythia8Generator)
PXL_PLUGIN_INIT