GeneratorProducts/test/rootio_t.cc

0001 #include "TFile.h"
0002 #include "TTree.h"
0003
0004 #include "HepMC/GenEvent.h"
0005 #include <cassert>
0006
0007 int main() {
0008   //Write
0009
0010   constexpr int eventNumber = 10;
0011   {
0012     TFile oFile("rootio_t.root", "RECREATE");
0013     TTree events("Events","Events");
0014
0015     HepMC::GenEvent event;
0016
0017     auto vtx = new HepMC::GenVertex( HepMC::FourVector(0.,0.,0.));
0018
0019     constexpr double kMass = 1.;
0020     constexpr double kMom = 1.;
0021     const double kEnergy = sqrt( kMass*kMass+kMom*kMom);
0022     constexpr int id = 11;
0023
0024     auto p1 = new HepMC::GenParticle(HepMC::FourVector{kMom,0.,0.,kEnergy}, id, 1);
0025
0026     p1->suggest_barcode(1);
0027     vtx->add_particle_out(p1);
0028
0029
0030     auto p2 = new HepMC::GenParticle(HepMC::FourVector{-kMom,0.,0.,kEnergy}, -id, 1);
0031
0032     p2->suggest_barcode(2);
0033     vtx->add_particle_out(p2);
0034
0035     auto decayVtx = new HepMC::GenVertex(HepMC::FourVector(1.,0.,0.));
0036     decayVtx->add_particle_in(p1);
0037
0038     event.add_vertex(vtx);
0039     event.add_vertex(decayVtx);
0040
0041     event.set_event_number(eventNumber);
0042     event.set_signal_process_id(20);
0043
0044     HepMC::GenEvent* pEv = &event;
0045
0046     events.Branch("GenEvent",&pEv);
0047
0048     events.Fill();
0049
0050     oFile.Write();
0051     oFile.Close();
0052   }
0053
0054   //Read
0055   {
0056     TFile iFile("rootio_t.root");
0057
0058     TTree* events = dynamic_cast<TTree*>( iFile.Get("Events") );
0059
0060     HepMC::GenEvent event;
0061     HepMC::GenEvent* pEv = &event;
0062     events->SetBranchAddress("GenEvent",&pEv);
0063
0064     events->GetEntry(0);
0065
0066     assert(event.event_number() == eventNumber);
0067     assert(event.particles_size() == 2);
0068     assert(event.vertices_size() == 2);
0069
0070     int barcode = 1;
0071     for(auto it = event.particles_begin(); it != event.particles_end(); ++it) {
0072       if((*it)->production_vertex()) {
0073         auto vtx = (*it)->production_vertex();
0074         assert( std::find(vtx->particles_out_const_begin(), vtx->particles_out_const_end(), *it) != vtx->particles_out_const_end());
0075       }
0076       if((*it)->end_vertex()) {
0077         auto vtx = (*it)->end_vertex();
0078         assert( std::find(vtx->particles_in_const_begin(), vtx->particles_in_const_end(), *it) != vtx->particles_in_const_end());
0079       }
0080
0081       assert( (*it)->barcode() == barcode);
0082       assert( *it == event.barcode_to_particle( (*it)->barcode() ) );
0083       ++barcode;
0084     }
0085
0086
0087   }
0088
0089 }