Last active
June 20, 2024 15:17
-
-
Save YanzhaoW/227c98d9de66c554c05b67ca5c6b6c6d to your computer and use it in GitHub Desktop.
ROOT macro script for the event stitching of neuland and sofia data.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#include "EventStitchingUtils.hpp" | |
void event_stitching() | |
{ | |
// Input set with neuland with a tree name | |
auto neuland_files = InputRootFiles{ "evt" }; | |
// Add filenames one by one. To make sure the algorithm work, always add the root file with smaller event IDs first | |
neuland_files.add_filename("/lustre/land/desha/s455/hit_plus_cluster0360_GTO.root"); | |
// neuland_files.add_filename("/lustre/land/desha/s455/hit_plus_cluster0360_GTO_1.root"); | |
neuland_files.init(); | |
// Input set for sofia with a tree name | |
auto sofia_files = InputRootFiles{ "ZcalIt7" }; | |
// Same above | |
sofia_files.add_filename("/lustre/land/desha/s455/s455_U238_ZcalIt7_dUmlenk_1730mm_HighGainR348.root"); | |
sofia_files.init(); | |
// Output with a tree name and a file name | |
auto output_file = OutputRootFile{ "evt", "test.root" }; | |
// Setup printing per certain number of events. | |
output_file.set_print_event_number(200); | |
// Output will write EventHeader branch from neuland file | |
output_file.record_branch("EventHeader.", neuland_files); | |
// Output will write NeulandHits branch from neuland file. The last value is the minimum size. The event will not be | |
// written if the size of hits is below the value. | |
output_file.record_tca_branch("NeulandHits", neuland_files, 3); | |
// Output will write NeulandClusters branch from neuland file. Same above. | |
output_file.record_tca_branch("NeulandClusters", neuland_files, 2); | |
// Output will write everything from sofia files | |
output_file.record_all_branches(sofia_files); | |
// Begin stitching the events from neuland and sofia | |
stitch_event(sofia_files, neuland_files, output_file); | |
// Write the output to a root file | |
output_file.write(); | |
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
namespace ROOT::Internal::RDF | |
{ | |
std::string GetBranchOrLeafTypeName(TTree& t, const std::string& colName); | |
} | |
class InputRootFiles | |
{ | |
public: | |
InputRootFiles(std::string_view tree_name) { data_ = std::make_unique<TChain>(tree_name.data()); } | |
void add_filename(std::string_view filename) | |
{ | |
data_->AddFile(filename.data(), TTree::kMaxEntries); | |
data_->LoadTree(TTree::kMaxEntries - 1); | |
} | |
auto init() -> TTreeReader& | |
{ | |
data_reader_.SetTree(data_.get()); | |
return data_reader_; | |
} | |
template <typename BranchType> | |
auto get_value_branch(std::string_view branch_name) | |
{ | |
return TTreeReaderValue<BranchType>{ data_reader_, branch_name.data() }; | |
} | |
template <typename BranchType> | |
auto get_array_branch(std::string_view branch_name) | |
{ | |
return TTreeReaderArray<BranchType>{ data_reader_, branch_name.data() }; | |
} | |
auto get_reader() -> TTreeReader& { return data_reader_; } | |
auto get_tree() -> TTree* { return data_.get(); } | |
private: | |
std::unique_ptr<TChain> data_; | |
TTreeReader data_reader_; | |
}; | |
class OutputRootFile | |
{ | |
public: | |
OutputRootFile(std::string_view tree_name, std::string_view filename) | |
: output_{ std::make_unique<TFile>(filename.data(), "recreate") } | |
, data_tree_{ tree_name.data(), "merged tree" } | |
{ | |
} | |
template <typename BranchType = TObject> | |
void record_branch(std::string_view branch_name, InputRootFiles& input) | |
{ | |
auto branch = std::make_unique<TTreeReaderValue<BranchType>>(input.get_reader(), branch_name.data()); | |
branch_setters_.push_back([this, &input, branch_name, branch = branch.get()]() mutable | |
{ data_tree_.Branch(branch_name.data(), branch->Get()); }); | |
branch_checkers_.push_back( | |
[this, branch = branch.get()]() | |
{ | |
if (branch->Get() != nullptr and branch->IsValid()) | |
{ | |
return true; | |
} | |
std::cerr << "Branch " << branch->GetBranchName() << " cannot be read!\n"; | |
return false; | |
}); | |
readers_.push_back(std::move(branch)); | |
} | |
void record_tca_branch(std::string_view branch_name, InputRootFiles& input, int min_size = 0) | |
{ | |
auto branch = std::make_unique<TTreeReaderValue<TClonesArray>>(input.get_reader(), branch_name.data()); | |
branch_setters_.push_back([this, &input, branch_name, branch = branch.get()]() mutable | |
{ data_tree_.Branch(branch_name.data(), "TClonesArray", branch->Get()); }); | |
branch_checkers_.push_back([this, branch = branch.get(), min_size]() mutable | |
{ return branch->Get()->GetEntries() > min_size; }); | |
readers_.push_back(std::move(branch)); | |
} | |
void record_all_branches(InputRootFiles& input) | |
{ | |
auto* tree_data = input.get_tree(); | |
auto* list_of_branches = tree_data->GetListOfBranches(); | |
for (auto* branch : TRangeDynCast<TBranch>(list_of_branches)) | |
{ | |
const auto* branch_name = branch->GetName(); | |
auto branch_type_name = ROOT::Internal::RDF::GetBranchOrLeafTypeName(*tree_data, branch_name); | |
if (branch_type_name.empty()) | |
{ | |
std::cerr << "Edwin: Cannot resolve the type of the branch " << branch_name << "\n"; | |
continue; | |
} | |
if (branch_type_name == "TClonesArray") | |
{ | |
record_tca_branch(branch_name, input); | |
} | |
else if (is_inherited_from_TObject(branch_type_name)) | |
{ | |
record_branch<TObject>(branch_name, input); | |
} | |
else if (branch_type_name == "Double_t") | |
{ | |
record_branch<double>(branch_name, input); | |
} | |
else if (branch_type_name == "Float_t") | |
{ | |
record_branch<float>(branch_name, input); | |
} | |
else if (branch_type_name == "UShort_t") | |
{ | |
record_branch<unsigned short>(branch_name, input); | |
} | |
else if (branch_type_name == "Short_t") | |
{ | |
record_branch<short>(branch_name, input); | |
} | |
else | |
{ | |
std::cerr << "Edwin: Type name " << branch_type_name << " is not supported!\n"; | |
} | |
} | |
} | |
auto fill() -> bool | |
{ | |
auto is_passed = true; | |
for (const auto& checker : branch_checkers_) | |
{ | |
is_passed &= checker(); | |
} | |
if (is_passed and data_tree_.Fill()) | |
{ | |
++entry_count; | |
if (entry_count % event_print_number == 0) | |
{ | |
std::cout << "Total number of the output events: " << entry_count << "\n"; | |
} | |
} | |
return is_passed; | |
} | |
void init() | |
{ | |
for (const auto& setter : branch_setters_) | |
{ | |
setter(); | |
} | |
} | |
void write() | |
{ | |
std::cout << "Writing output file ... \n"; | |
output_->WriteObject<TTree>(&data_tree_, data_tree_.GetName()); | |
std::cout << "Output file has been written \n"; | |
} | |
void set_print_event_number(unsigned int num) { event_print_number = num; } | |
private: | |
unsigned int event_print_number = 100; | |
std::unique_ptr<TFile> output_; | |
std::vector<std::unique_ptr<ROOT::Internal::TTreeReaderValueBase>> readers_; | |
std::vector<std::function<void()>> branch_setters_; | |
std::vector<std::function<bool()>> branch_checkers_; | |
TTree data_tree_; | |
uint64_t entry_count = 0; | |
auto is_inherited_from_TObject(std::string_view type_name) -> bool | |
{ | |
auto* type_class = TClass::GetClass(type_name.data()); | |
return type_class != nullptr and type_class->InheritsFrom("TObject"); | |
} | |
}; | |
void stitch_event(InputRootFiles& sofia_files, InputRootFiles& r3b_files, OutputRootFile& output) | |
{ | |
std::cout << "Start Event stitching ...\n"; | |
auto& r3b_reader = r3b_files.get_reader(); | |
auto& sofia_reader = sofia_files.get_reader(); | |
auto event_header = r3b_files.get_value_branch<R3BEventHeader>("EventHeader."); | |
auto sofia_event_id = sofia_files.get_value_branch<UShort_t>("fEventno"); | |
auto is_not_finished = true; | |
is_not_finished &= sofia_reader.Next(); | |
is_not_finished &= r3b_reader.Next(); | |
output.init(); | |
std::cout << "Starting to loop through events...\n"; | |
while (is_not_finished) | |
{ | |
const auto r3b_event_number = static_cast<UShort_t>(event_header->GetEventno()); | |
const auto sofia_event_number = *sofia_event_id; | |
if (r3b_event_number == sofia_event_number) | |
{ | |
output.fill(); | |
} | |
if (r3b_event_number >= sofia_event_number) | |
{ | |
is_not_finished &= sofia_reader.Next(); | |
} | |
if (r3b_event_number <= sofia_event_number) | |
{ | |
is_not_finished &= r3b_reader.Next(); | |
} | |
} | |
std::cout << "Event stitching finished\n"; | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
How to use
Download these two files into the same folder. In
event_stitching.cpp
, change the input file names and other information according to your needs.To run the program:
Contact me if you have any problems or suggestions.