File size: 8,586 Bytes
cfcbbc8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 |
#!/usr/bin/env python3
"""
Simple ROOT file inspector using uproot
Usage: python inspect_root.py [options]
Default data directory: /global/cfs/projectdirs/atlas/eligd/llm_for_analysis_copy/data/
"""
import sys
import os
import uproot
import numpy as np
import argparse
# Default data directory
DEFAULT_DATA_DIR = '/global/cfs/projectdirs/atlas/eligd/llm_for_analysis_copy/data/'
def inspect_root_file(filepath, show_stats=True, max_entries=None):
"""Inspect a ROOT file and display its contents"""
if not os.path.exists(filepath):
print("Error: File '{}' does not exist!".format(filepath))
return
print("Inspecting ROOT file: {}".format(filepath))
print("=" * 80)
try:
# Open the ROOT file
with uproot.open(filepath) as file:
print("File opened successfully!")
print("Keys in file: {}".format(list(file.keys())))
print()
# Look for trees (typically named 'mini' or similar)
trees = [key for key in file.keys() if not key.endswith(';1') or ';' not in key]
print("Available trees/objects: {}".format(trees))
print()
# Try to find and inspect the main tree
tree_name = None
if 'mini;1' in file:
tree_name = 'mini;1'
elif 'mini' in file:
tree_name = 'mini'
else:
# Look for the first tree-like object
for key in file.keys():
if ';' in key and not key.endswith('.root'):
tree_name = key
break
if tree_name:
print("Inspecting tree: {}".format(tree_name))
tree = file[tree_name]
# Get basic tree information
num_entries = tree.num_entries
print("Number of entries: {}".format(num_entries))
print("Branches: {}".format(list(tree.keys())))
print()
if max_entries:
num_entries = min(num_entries, max_entries)
print("Limiting analysis to first {} entries".format(num_entries))
print()
# Show branch details
print("Branch details:")
for branch_name in tree.keys():
try:
branch = tree[branch_name]
if hasattr(branch, 'array'):
arr = branch.array()
if max_entries:
arr = arr[:max_entries]
# Handle different array shapes
if hasattr(arr, 'shape'):
print(" {}: shape={}, dtype={}".format(branch_name, arr.shape, arr.dtype))
else:
# Handle scalar or other types
print(" {}: type={}, value_type={}".format(branch_name, type(arr), arr.dtype if hasattr(arr, 'dtype') else 'unknown'))
if len(arr) > 0:
sample_size = min(5, len(arr))
print(" Sample values: {}".format(arr[:sample_size]))
else:
print(" {}: {}".format(branch_name, type(branch)))
except Exception as e:
print(" {}: Error reading - {}".format(branch_name, e))
print()
if show_stats:
# Show some basic statistics for numerical branches
print("Basic statistics for numerical branches:")
for branch_name in tree.keys():
try:
branch = tree[branch_name]
if hasattr(branch, 'array'):
arr = branch.array()
if max_entries:
arr = arr[:max_entries]
if arr.dtype.kind in 'iufc': # integer, unsigned, float, complex
print(" {}:".format(branch_name))
print(" Mean: {:.3f}".format(np.mean(arr)))
print(" Std: {:.3f}".format(np.std(arr)))
print(" Min: {:.3f}".format(np.min(arr)))
print(" Max: {:.3f}".format(np.max(arr)))
except:
pass
else:
print("No tree found in the file.")
except Exception as e:
print("Error opening file: {}".format(e))
def list_files_in_directory(directory):
"""List all ROOT files in a directory"""
if not os.path.exists(directory):
print("Directory '{}' does not exist!".format(directory))
return []
root_files = []
for filename in os.listdir(directory):
if filename.endswith('.root'):
root_files.append(os.path.join(directory, filename))
return sorted(root_files)
def main():
parser = argparse.ArgumentParser(description='Inspect ROOT files using uproot')
parser.add_argument('filepath', nargs='?', help='Path to the ROOT file to inspect')
parser.add_argument('--data-dir', default=DEFAULT_DATA_DIR,
help='Default data directory (default: {})'.format(DEFAULT_DATA_DIR))
parser.add_argument('--list-files', action='store_true',
help='List all ROOT files in the data directory')
parser.add_argument('--inspect-all', action='store_true',
help='Inspect all ROOT files in the data directory')
parser.add_argument('--no-stats', action='store_true',
help='Skip statistical analysis')
parser.add_argument('--max-entries', type=int,
help='Limit analysis to first N entries')
args = parser.parse_args()
if args.list_files:
# Just list files
root_files = list_files_in_directory(args.data_dir)
print("ROOT files in {}:".format(args.data_dir))
for i, filepath in enumerate(root_files, 1):
filename = os.path.basename(filepath)
print(" {}. {}".format(i, filename))
return
if args.inspect_all:
# Inspect all files
root_files = list_files_in_directory(args.data_dir)
if not root_files:
print("No ROOT files found in {}".format(args.data_dir))
return
print("Inspecting all {} ROOT files in {}".format(len(root_files), args.data_dir))
print("=" * 80)
for i, filepath in enumerate(root_files, 1):
filename = os.path.basename(filepath)
print("\n[{}/{}] Inspecting: {}".format(i, len(root_files), filename))
print("-" * 40)
inspect_root_file(filepath, show_stats=not args.no_stats, max_entries=args.max_entries)
return
# Single file inspection
if not args.filepath:
# If no filepath provided, show available files and prompt
root_files = list_files_in_directory(args.data_dir)
if root_files:
print("Available ROOT files in {}:".format(args.data_dir))
for i, filepath in enumerate(root_files, 1):
filename = os.path.basename(filepath)
print(" {}. {}".format(i, filename))
try:
choice = input("\nEnter file number or full path (or 'q' to quit): ").strip()
if choice.lower() == 'q':
return
if choice.isdigit():
idx = int(choice) - 1
if 0 <= idx < len(root_files):
filepath = root_files[idx]
else:
print("Invalid choice!")
return
else:
filepath = choice
except KeyboardInterrupt:
print("\nExiting...")
return
else:
filepath = input("Enter path to ROOT file: ").strip()
else:
filepath = args.filepath
# Handle relative paths
if not os.path.isabs(filepath):
filepath = os.path.join(args.data_dir, filepath)
inspect_root_file(filepath, show_stats=not args.no_stats, max_entries=args.max_entries)
if __name__ == "__main__":
main()
|