src/root/JetScapeWriterRootHepMC.h

0001 /*******************************************************************************
0002  * Copyright (c) The JETSCAPE Collaboration, 2018
0003  *
0004  * Modular, task-based framework for simulating all aspects of heavy-ion collisions
0005  *
0006  * For the list of contributors see AUTHORS.
0007  *
0008  * Report issues at https://github.com/JETSCAPE/JETSCAPE/issues
0009  *
0010  * or via email to bugs.jetscape@gmail.com
0011  *
0012  * Distributed under the GNU General Public License 3.0 (GPLv3 or later).
0013  * See COPYING for details.
0014  ******************************************************************************/
0015
0016 #ifndef JETSCAPEWRITERROOTHEPMC_H
0017 #define JETSCAPEWRITERROOTHEPMC_H
0018
0019 #include <fstream>
0020 #include <string>
0021
0022 #include "JetScapeWriter.h"
0023 #include "PartonShower.h"
0024
0025 #include "HepMC3/GenEvent.h"
0026 #include "HepMC3/ReaderAscii.h"
0027 #include "HepMC3/WriterRootTree.h"
0028 #include "HepMC3/Print.h"
0029
0030 // using namespace HepMC;
0031 using HepMC3::GenEvent;
0032 using HepMC3::GenVertex;
0033 using HepMC3::GenParticle;
0034 using HepMC3::GenVertexPtr;
0035 using HepMC3::GenParticlePtr;
0036
0037 namespace Jetscape {
0038
0039 class JetScapeWriterRootHepMC : public JetScapeWriter, public HepMC3::WriterRootTree {
0040
0041 public:
0042   JetScapeWriterRootHepMC() : HepMC3::WriterRootTree("") { SetId("HepMC ROOT writer"); };
0043   JetScapeWriterRootHepMC(string m_file_name_out)
0044       : JetScapeWriter(m_file_name_out), HepMC3::WriterRootTree(m_file_name_out) {
0045     SetId("HepMC ROOT writer");
0046   };
0047   virtual ~JetScapeWriterRootHepMC();
0048
0049   void Init();
0050   void Exec();
0051
0052   bool GetStatus() { return failed(); }
0053   void Close() { close(); }
0054
0055   // // NEVER use this!
0056   // // Can work with only one writer, but with a second one it gets called twice
0057   // void WriteTask(weak_ptr<JetScapeWriter> w);
0058
0059   // overload write functions
0060   void WriteEvent();
0061
0062   // At parton level, we should never accept anything other than a full shower
0063   // void Write(weak_ptr<Vertex> v);
0064   void Write(weak_ptr<PartonShower> ps);
0065   void Write(weak_ptr<Hadron> h);
0066   void WriteHeaderToFile();
0067
0068 private:
0069   HepMC3::GenEvent evt;
0070   vector<HepMC3::GenVertexPtr> vertices;
0071   HepMC3::GenVertexPtr hadronizationvertex;
0072
0073   /// WriteEvent needs to know whether it should overwrite final partons status to 1
0074   bool hashadrons=false;
0075
0076   inline HepMC3::GenVertexPtr
0077   castVtxToHepMC(const shared_ptr<Vertex> vtx) const {
0078     double x = vtx->x_in().x();
0079     double y = vtx->x_in().y();
0080     double z = vtx->x_in().z();
0081     double t = vtx->x_in().t();
0082     HepMC3::FourVector vtxPosition(x, y, z, t);
0083     // if ( t< 1e-6 ) t = 1e-6; // could do this. Exact 0 is bit quirky but works for hepmc
0084     return make_shared<GenVertex>(vtxPosition);
0085   }
0086
0087   inline HepMC3::GenParticlePtr
0088   castPartonToHepMC(const shared_ptr<Parton> pparticle) const {
0089     return castPartonToHepMC(*pparticle);
0090   }
0091
0092   inline HepMC3::GenParticlePtr
0093   castPartonToHepMC(const Parton &particle) const {
0094     HepMC3::FourVector pmom(particle.px(), particle.py(), particle.pz(),
0095                             particle.e());
0096     return make_shared<GenParticle>(pmom, particle.pid(), particle.pstat());
0097   }
0098
0099   inline HepMC3::GenParticlePtr
0100   castHadronToHepMC(const shared_ptr<Hadron> pparticle) const {
0101     return castHadronToHepMC(*pparticle);
0102   }
0103
0104   inline HepMC3::GenParticlePtr
0105   castHadronToHepMC(const Hadron &particle) const {
0106     HepMC3::FourVector pmom(particle.px(), particle.py(), particle.pz(),
0107                             particle.e());
0108     return make_shared<GenParticle>(pmom, particle.pid(), particle.pstat());
0109   }
0110
0111   //int m_precision; //!< Output precision
0112 };
0113
0114 } // end namespace Jetscape
0115
0116 #endif