| 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 | {
|
|---|
| 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 |
|
|---|
| 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 |
|
|---|
| 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 |
|
|---|
| 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 |
|
|---|
| 91 | void WriterRootTree::write_run_info() {}
|
|---|
| 92 |
|
|---|
| 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 |
|
|---|
| 101 | bool WriterRootTree::failed()
|
|---|
| 102 | {
|
|---|
| 103 | if ( !m_file->IsOpen() ) return true;
|
|---|
| 104 |
|
|---|
| 105 | return false;
|
|---|
| 106 | }
|
|---|
| 107 |
|
|---|
| 108 | } // namespace HepMC3
|
|---|
