LLM4HEP / get_arr.py
ho22joshua's picture
initial commit
cfcbbc8
import numpy as np
import argparse
import os
parser = argparse.ArgumentParser(description='read array')
add_arg = parser.add_argument
add_arg('--name', help='array name')
add_arg('--out_dir', help='output directory', default='.')
args = parser.parse_args()
# Prefer arrays saved under <out_dir>/logs, fallback to current directory
logs_path = os.path.join(args.out_dir, 'logs', f'{args.name}.npy')
root_path = os.path.join(args.out_dir, f'{args.name}.npy')
filepath = logs_path if os.path.exists(logs_path) else root_path
arr = np.load(filepath)
if len(arr) > 3:
arr = np.array([np.sum(arr[:-2]), arr[-2], arr[-1]])
print(*arr.flatten())