1 // -*- C++ -*-
2 //
3 // This file is part of HepMC
4 // Copyright (C) 2014-2019 The HepMC collaboration (see AUTHORS for details)
5 //
6 /**
7 * @file WriterRootTree.cc
8 * @brief Implementation of \b class WriterRootTree
9 *
10 */
11 #include <cstdio> // sprintf
12 #include "HepMC3/WriterRootTree.h"
13 #include "HepMC3/Version.h"
14 // ROOT header files
15 #include "TFile.h"
16 #include "TTree.h"
17
18 namespace HepMC3
19 {
HEPMC3_DECLARE_WRITER_FILE(WriterRootTree)20 HEPMC3_DECLARE_WRITER_FILE(WriterRootTree)
21
22 WriterRootTree::WriterRootTree(const std::string &filename, std::shared_ptr<GenRunInfo> run):
23 m_tree(0),
24 m_events_count(0),
25 m_tree_name("hepmc3_tree"),
26 m_branch_name("hepmc3_event")
27 {
28 m_file = TFile::Open(filename.c_str(), "RECREATE");
29 if (!init(run)) return;
30 }
31
WriterRootTree(const std::string & filename,const std::string & treename,const std::string & branchname,std::shared_ptr<GenRunInfo> run)32 WriterRootTree::WriterRootTree(const std::string &filename, const std::string &treename, const std::string &branchname, std::shared_ptr<GenRunInfo> run):
33 m_tree(0),
34 m_events_count(0),
35 m_tree_name(treename.c_str()),
36 m_branch_name(branchname.c_str())
37 {
38 m_file = TFile::Open(filename.c_str(), "RECREATE");
39 if (!init(run)) return;
40 }
41
init(std::shared_ptr<GenRunInfo> run)42 bool WriterRootTree::init(std::shared_ptr<GenRunInfo> run )
43 {
44 if ( !m_file->IsOpen() )
45 {
46 HEPMC3_ERROR("WriterRootTree: problem opening file: " << m_file->GetName())
47 return false;
48 }
49 m_event_data = new GenEventData();
50 m_run_info_data = new GenRunInfoData();
51 set_run_info(run);
52 if ( run_info() ) run_info()->write_data(*m_run_info_data);
53 m_tree = new TTree(m_tree_name.c_str(), "hepmc3_tree");
54 m_tree->Branch(m_branch_name.c_str(), m_event_data);
55 m_tree->Branch("GenRunInfo", m_run_info_data);
56 return true;
57 }
58
write_event(const GenEvent & evt)59 void WriterRootTree::write_event(const GenEvent &evt)
60 {
61 if ( !m_file->IsOpen() ) return;
62 bool refill = false;
63 if ( evt.run_info()&&(!run_info() || (run_info() != evt.run_info()))) { set_run_info(evt.run_info()); refill = true;}
64 if (refill)
65 {
66 m_run_info_data->weight_names.clear();
67 m_run_info_data->tool_name.clear();
68 m_run_info_data->tool_version.clear();
69 m_run_info_data->tool_description.clear();
70 m_run_info_data->attribute_name.clear();
71 m_run_info_data->attribute_string.clear();
72 run_info()->write_data(*m_run_info_data);
73 }
74
75
76
77 m_event_data->particles.clear();
78 m_event_data->vertices.clear();
79 m_event_data->links1.clear();
80 m_event_data->links2.clear();
81 m_event_data->attribute_id.clear();
82 m_event_data->attribute_name.clear();
83 m_event_data->attribute_string.clear();
84
85 evt.write_data(*m_event_data);
86 m_tree->Fill();
87 ++m_events_count;
88 }
89
90
write_run_info()91 void WriterRootTree::write_run_info() {}
92
close()93 void WriterRootTree::close()
94 {
95 m_file->WriteTObject(m_tree);
96 m_file->Close();
97 delete m_event_data;
98 delete m_run_info_data;
99 }
100
failed()101 bool WriterRootTree::failed()
102 {
103 if ( !m_file->IsOpen() ) return true;
104
105 return false;
106 }
107
108 } // namespace HepMC3
109