pyrobird/tests/ttree_cutter.py

0001 #!/usr/bin/env python3
0002
0003 import ROOT
0004 import argparse
0005 import sys
0006
0007
0008 def copy_events(input_filename, output_filename, tree_name, num_events_to_copy):
0009     # Open the original ROOT file
0010     input_file = ROOT.TFile.Open(input_filename, "READ")
0011     if not input_file or input_file.IsZombie():
0012         print(f"Error: Cannot open input file '{input_filename}'.")
0013         sys.exit(1)
0014
0015     # Get the TTree from the original file
0016     input_tree = input_file.Get(tree_name)
0017     if not input_tree:
0018         print(f"Error: Tree '{tree_name}' not found in '{input_filename}'.")
0019         input_file.Close()
0020         sys.exit(1)
0021
0022     # Create a new ROOT file
0023     output_file = ROOT.TFile.Open(output_filename, "RECREATE")
0024     if not output_file or output_file.IsZombie():
0025         print(f"Error: Cannot create output file '{output_filename}'.")
0026         input_file.Close()
0027         sys.exit(1)
0028
0029     # Set maximum compression level
0030     output_file.SetCompressionLevel(9)
0031
0032     # Clone the tree structure (no events are copied yet)
0033     output_tree = input_tree.CloneTree(0)
0034
0035     # Total number of events in the input tree
0036     total_events = input_tree.GetEntries()
0037     if num_events_to_copy > total_events:
0038         print(f"Warning: Requested {num_events_to_copy} events, but the input tree only contains {total_events} events.")
0039         num_events_to_copy = total_events
0040
0041     # Copy the desired events
0042     for i in range(num_events_to_copy):
0043         input_tree.GetEntry(i)
0044         output_tree.Fill()
0045
0046     # Write the new tree to the output file
0047     output_tree.Write()
0048
0049     # Close the files
0050     output_file.Close()
0051     input_file.Close()
0052
0053     print(f"Successfully created '{output_filename}' with {num_events_to_copy} event(s).")
0054
0055
0056 def main():
0057     parser = argparse.ArgumentParser(description='Copy a specified number of events from a ROOT file to a new ROOT file.')
0058     parser.add_argument('input_file', help='Path to the input ROOT file.')
0059     parser.add_argument('output_file', help='Path to the output ROOT file.')
0060     parser.add_argument('-n', '--num-events', type=int, default=2, help='Number of events to copy (default: 2).')
0061     parser.add_argument('-t', '--tree-name', default='events', help='Name of the TTree to copy from (default: myTree).')
0062
0063     args = parser.parse_args()
0064
0065     copy_events(args.input_file, args.output_file, args.tree_name, args.num_events)
0066
0067
0068 if __name__ == '__main__':
0069     main()