source: trunk/Mars/hawc/stripandreco.py@ 20094

import os
import argparse
import subprocess
import time

if (os.environ.get('CONFIG_HAWC') == None):
    print "Did not find HAWC enviroment (try hawc-sw)!"
    exit()

parser = argparse.ArgumentParser()
parser.add_argument("-p", "--path", dest="path", help="Search path (recursive!) for trigger files (xcd)",type=str,required=True)
parser.add_argument("-s", "--strip", action="store_true", help="Only strip all files in path", required=False)
parser.add_argument("-r", "--reco", action="store_true", help="Only reconstruct all files in path", required=False)
parser.add_argument("-c", "--root", action="store_true", help="Only convert all files to root in path", required=False)
parser.add_argument("-j", "--proc", dest="proc", help="Number of processes (max. cores + 1)",type=int , required=False)
parser.add_argument("--filter", dest="filter", help="Only use files containing string FILTER", type=str, required=False)
parser.add_argument("--channel", dest="channel", help="Specify a channel for the trigger flag (Default: 31,E16C)",type=int , required=False)

args = parser.parse_args()

channel_id = 31
max_proc = 2

if args.channel:
    channel_id = args.channel

if args.proc:
    max_proc = args.proc

print """
************************************
*                                  *
*       Welcome to HAWCs Eye       *
*                                  *
*  Filter and reconstruction tool  *
*                                  *
*     Version 1.0 01.08.2017       *
*                                  *
************************************

Search path: {}
Channel ID : {}
""".format(args.path,channel_id)

process_list = []
if not args.reco and not args.root:
    strip_flist = []
    #Search for .xcd files in args.path without "strp" or "reco" in the filename
    for root,dirs,files in os.walk(args.path):
        if args.filter is None:
            for xcdfile in [xcdf for xcdf in files if (('.xcd' in xcdf[-4:]) and not ('strp' in xcdf) and not ('reco' in xcdf))]:
                strip_flist.append(root+"/"+xcdfile)
        else:
            for xcdfile in [xcdf for xcdf in files if (('.xcd' in xcdf[-4:]) and not ('strp' in xcdf) and not ('reco' in xcdf) and (args.filter in xcdf))]:
                strip_flist.append(root+"/"+xcdfile)


    #Filter files for a hit in channel_id
    for xcdfile in strip_flist:
        process_string = 'xcdf select "trig.channelID_4Edge == {}" -o {}_strp.xcd {}'.format(channel_id,xcdfile[:-4],xcdfile)
        print "Starting: " + process_string
        process_list.append(subprocess.Popen(process_string,shell=True))
        while([status.poll() for status in process_list].count(None) >= max_proc):
            time.sleep(0.2)

#Wait for all strip processes to finish:
while([status.poll() for status in process_list].count(None) > 0):
    time.sleep(0.2)
#Reset process list:
process_list = []

if not args.strip and not args.root:
    reco_flist = []
    #Search for .xcd files in args.path with "strp" in the filename
    for root,dirs,files in os.walk(args.path):
        if args.filter is None:
            for xcdfile in [xcdf for xcdf in files if (('.xcd' in xcdf[-4:]) and ('strp' in xcdf[-8:-4]))]:
                reco_flist.append(root+"/"+xcdfile)
        else:
            for xcdfile in [xcdf for xcdf in files if (('.xcd' in xcdf[-4:]) and ('strp' in xcdf[-8:-4]) and (args.filter in xcdf))]:
                reco_flist.append(root+"/"+xcdfile)

    #Reco strp files:
    for xcdfile in reco_flist:
        if (args.filter == None):
            process_string = 'offline-reconstructor --input {} -o {}_reco.xcd'.format(xcdfile,xcdfile[:-9])
        else:
            process_string = 'offline-reconstructor --input {} -o {}_reco.xcd'.format(xcdfile,xcdfile[:-4])
        print "Starting: " + process_string
        process_list.append(subprocess.Popen(process_string,shell=True))
        while([status.poll() for status in process_list].count(None) >= max_proc):
            time.sleep(0.2)

#Wait for all reconstruction processes to finish
while([status.poll() for status in process_list].count(None) > 0):
    time.sleep(0.2)

#Reset process list:
process_list = []

if not args.strip and not args.reco:
    root_flist = []
    #Search for .xcd files in args.path with "reco" in the filename
    for root,dirs,files in os.walk(args.path):
        if args.filter is None:
            for xcdfile in [xcdf for xcdf in files if (('.xcd' in xcdf[-4:]) and ('reco' in xcdf))]:
                root_flist.append(root+"/"+xcdfile)
        else:
            for xcdfile in [xcdf for xcdf in files if (('.xcd' in xcdf[-4:]) and ('reco' in xcdf) and (args.filter in xcdf))]:
                root_flist.append(root+"/"+xcdfile)

    #xcd to root files:
    for xcdfile in root_flist:
        process_string = 'xcdf-root --input {} -o {}.root'.format(xcdfile,xcdfile[:-4])
        print "Starting: " + process_string
        process_list.append(subprocess.Popen(process_string,shell=True))
        while([status.poll() for status in process_list].count(None) >= max_proc):
            time.sleep(0.2)

#Wait for all files beeing converted
while([status.poll() for status in process_list].count(None) > 0):
    time.sleep(0.2)