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     auto p2 = new HepMC::GenParticle(HepMC::FourVector{-kMom, 0., 0., kEnergy}, -id, 1);
0030
0031     p2->suggest_barcode(2);
0032     vtx->add_particle_out(p2);
0033
0034     auto decayVtx = new HepMC::GenVertex(HepMC::FourVector(1., 0., 0.));
0035     decayVtx->add_particle_in(p1);
0036
0037     event.add_vertex(vtx);
0038     event.add_vertex(decayVtx);
0039
0040     event.set_event_number(eventNumber);
0041     event.set_signal_process_id(20);
0042
0043     HepMC::GenEvent* pEv = &event;
0044
0045     events.Branch("GenEvent", &pEv);
0046
0047     events.Fill();
0048
0049     oFile.Write();
0050     oFile.Close();
0051   }
0052
0053   //Read
0054   {
0055     TFile iFile("rootio_t.root");
0056
0057     TTree* events = dynamic_cast<TTree*>(iFile.Get("Events"));
0058
0059     HepMC::GenEvent event;
0060     HepMC::GenEvent* pEv = &event;
0061     events->SetBranchAddress("GenEvent", &pEv);
0062
0063     events->GetEntry(0);
0064
0065     assert(event.event_number() == eventNumber);
0066     assert(event.particles_size() == 2);
0067     assert(event.vertices_size() == 2);
0068
0069     int barcode = 1;
0070     for (auto it = event.particles_begin(); it != event.particles_end(); ++it) {
0071       if ((*it)->production_vertex()) {
0072         auto vtx = (*it)->production_vertex();
0073         assert(std::find(vtx->particles_out_const_begin(), vtx->particles_out_const_end(), *it) !=
0074                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) !=
0079                vtx->particles_in_const_end());
0080       }
0081
0082       assert((*it)->barcode() == barcode);
0083       assert(*it == event.barcode_to_particle((*it)->barcode()));
0084       ++barcode;
0085     }
0086   }
0087 }