makent

#!/usr/bin/env python3
import os
import sys
import argparse

cli = argparse.ArgumentParser(description='Convert a text file to a ROOT ntuple.', epilog='If column names are not specified via the command-line, the first row of the text file is used if it contains strings, else generic, numbered names are used.')
cli.add_argument('-c', help='colon-separated column names')
cli.add_argument('-o', help='output ROOT file')
cli.add_argument('file', help='input TXT file', type=str)
args = cli.parse_args(sys.argv[1:])

if args.o is None:
    args.o = './%s.root'%os.path.basename(args.file)
if not os.path.isfile(args.file):
    cli.error('Input file does not exist:  '+args.file)
if os.path.isfile(args.o):
    cli.error('Output file already exists:  '+args.o)

import ROOT
import numpy

with open(args.file,'r') as file:

    for line in file.readlines():
        first_row = line.strip().split()
        break

    try:
        [ float(x) for x in first_row ]
        header = False
    except Exception as e:
        header = True

    if args.c is None:
        if header:
            args.c = ':'.join(first_row)
        else:
            args.c = ':'.join( [ 'x%d'%x for x in range(len(first_row)) ] )

    elif len(args.c.split(':')) != len(first_row):
        cli.error('Column count mismatch between -c and first row in file.')

    file.seek(0)
    g = ROOT.TFile(args.o, 'CREATE')
    t = ROOT.TNtupleD('t', 'makent', args.c)
    for n,line in enumerate(file.readlines()):
      if header and n == 0:
          continue
      if line.startswith('#'):
          continue
      try:
          x = [ float(y) for y in line.strip().split() ]
          t.Fill(numpy.array(x))
      except Exception as e:
          print('ERROR:  Ignoring bad data on row #%d:  %s'%(t.GetEntries(),line.strip()))
    g.Write()

print('Created %s with %d entries from %s with column names %s.'%(args.o,t.GetEntries(),args.file,args.c))