merge_histograms.py

Go to the documentation of this file.

##################################################
##
## merge_histograms.py
##
## Laura Fields
## Jul 2014
##
## Merges many flux histogram files into one
##
###################################################

from optparse import OptionParser
import sys, os, subprocess, shutil, ROOT,datetime

def is_good_histo_file(myfile):
    # Check that file is not a zombie
    if myfile.IsZombie():
        return False
    # Check that numu flux is not empty
    keys = myfile.GetListOfKeys()
    if "numu_flux" in keys:
        if myfile.Get("numu_flux").Integral()>0:
            return True
    if "numu_flux_forplots" in keys:
        if myfile.Get("numu_flux_forplots").Integral()>0:
            return True
    return False

def merge_histos(filevec,outfile):

    if len(filevec)==0:
        print "merge_hitsograms.py, mergehistos(): No files to merge."
        return;

    # Temporary staging area
    #temppath = "/dune/data2/users/"+os.getenv("USER")+"/scratch/"
    #if not os.path.exists(temppath):
    #    os.makedirs(temppath)

    #for i in range (0,len(filevec)):
    #    tempfile = temppath + "/" +os.path.basename(filevec[i])
    #    shutil.copyfile(filevec[i],tempfile)

    # find first good file
    first_file_index = -1
    
    for i in range(0,len(filevec)):
        #tempfile = temppath + "/" +os.path.basename(filevec[i])
        in_tfile = ROOT.TFile(filevec[i]);
        if is_good_histo_file(in_tfile):
            first_file_index = i
            break

    n_files = 1

    print "First good file: "+str(first_file_index)

    last_time = datetime.datetime.now() 
    # Look through first file to find list of histograms that we should merge
    output_histos = []
    for key in in_tfile.GetListOfKeys():
        if key.ReadObj().IsA().InheritsFrom("TH1"):
            output_histos.append(in_tfile.Get(key.GetName()).Clone())

    for i in range(first_file_index+1,len(filevec)):
        if i%50 == 0: 
            now = datetime.datetime.now()
            print "Merging file number",i,"; time to merge last 50 files:",now-last_time
            last_time = now
            
        #tempfile = temppath + "/" +os.path.basename(filevec[i])
        next_tfile = ROOT.TFile(filevec[i])
        if is_good_histo_file(next_tfile):
            n_files = n_files + 1
            for hist in output_histos:
                temphist = next_tfile.Get(hist.GetName())
                hist.Add(temphist)
        next_tfile.Close()
        #os.remove(tempfile)
    out_tfile = ROOT.TFile(outfile,"RECREATE");
    print "writing to "+outfile
    scale_factor = 1.0/n_files;
    for hist in output_histos:
        hist.Scale(scale_factor)
        hist.Write()
        
    return


def merge_globes(filevec,outfile):

    if len(filevec)==0:
        print "merge_histograms.py, merge_globes(): No files to merge."
        return;

    n_files = 0
            
    total_flux = []
    for i in range(501):
        total_flux.append([0,0,0,0,0,0,0])

    #temppath = "/scratch/dune/"+os.getenv("USER")
    #if not os.path.exists(temppath):
    #    os.makedirs(temppath)

    print "merging",len(filevec),"globes files"
    last_time = datetime.datetime.now() 
    for i in range(0,len(filevec)):
        if i%50==0: 
            now = datetime.datetime.now()
            print "Merging file number",i,"; time since last print:",now-last_time
            last_time = now

        #tempfile = temppath+"/"+os.path.basename(filevec[i])
        #shutil.copyfile(filevec[i],tempfile)
        next_file = open(filevec[i])
        lines = next_file.readlines()
        next_file.close()
        #os.remove(tempfile)

        flux = []
        problem_found = (len(lines)!=501)
        for line in lines:

            if len(line.split())!=7: 
                problem_found = True
            elif float(line.split()[0])>1 and float(line.split()[0])<5 and (float(line.split()[1])+float(line.split()[2])+float(line.split()[3])+float(line.split()[4])+float(line.split()[5])+float(line.split()[6]))==0.0:
                # make sure fluxes aren't all zero (by looking in peak)
                    problem_found = True
            else:
                flux.append([float(line.split()[0]),float(line.split()[1]),float(line.split()[2]),float(line.split()[3]),float(line.split()[4]),float(line.split()[5]),float(line.split()[6])])
    
        if not problem_found:
            n_files = n_files+1       
            for j in range(501):
                total_flux[j][0] = flux[j][0]
                for k in range(1,7):
                    total_flux[j][k] += flux[j][k]
            
    print "Number of good globes files found:",n_files
    for i in range(501):
        for j in range(1,7):
            total_flux[i][j] = total_flux[i][j]/n_files

    out_file = open(outfile,"w");
    print "writign to "+outfile
    for i in range(501):
        for j in range(7):
            out_file.write(str(total_flux[i][j])+" ")
        out_file.write("\n")

    return



###################################################
#
# Determine default g4lbne directory
# (the directory this script is in)
#
###################################################
scriptdir = os.path.abspath(sys.argv[0]+"/../..")
print "default g4lbne directory:",scriptdir

###################################################
#
# Setup parser that reads in options
#
###################################################

usage = "usage: %prog [options]"
parser = OptionParser(usage=usage)

parser.add_option("-g", "--g4lbne_dir", dest="g4lbne_dir",
                  help="g4lbne directory", default=scriptdir)
parser.add_option("-p", "--physics_list", dest="physics_list",
                  help="Geant4 Physics List", default="QGSP_BERT");
parser.add_option("-m", "--macro", dest="macro",
                  help="Template Macro", default="Nominal")
parser.add_option("-d", "--mode", dest="mode",
                  help="neutrino or antineutrino mode", default="neutrino")
parser.add_option("-n", "--n_events", dest="n_events",
                  help="Number of events per job", default=100000)
parser.add_option("-i", "--input_dir", dest="input_dir",
                  help="Input directory", default="/pnfs/dune/scratch/users/$USER/fluxfiles/g4lbne")
parser.add_option("-o", "--output_dir", dest="output_dir",
                  help="Input directory", default="/dune/data/users/$USER/fluxfiles/g4lbne")
parser.add_option("-l","--location",dest="location",help="detector location",default="LBNEFD")
parser.add_option("-v","--version",dest="version",help="g4lbne version number",default="v3r5p4")

(options, args) = parser.parse_args()


if options.g4lbne_dir.find("$USER")>=0:
    options.g4lbne_dir = options.g4lbne_dir.replace("$USER",os.getenv("USER"));
if options.input_dir.find("$USER")>=0:
    options.input_dir = options.input_dir.replace("$USER",os.getenv("USER"));
if options.output_dir.find("$USER")>=0:
    options.output_dir = options.output_dir.replace("$USER",os.getenv("USER"));

physics_list = options.physics_list
print "Using GEANT4 physics list",physics_list

###################################################
#
# Determine G4LBNE Version
#
###################################################
version=options.version

###################################################
#
# Print options
#
###################################################
print "Using the g4lbne installed at",options.g4lbne_dir
print "Assuming each file corresponds to ",options.n_events,"protons on target."
print "Using macro "+options.macro
print "Using input location:",options.input_dir
print "Using output location:",options.output_dir
print "Using location:",options.location

###################################################
#
# Find files
#
###################################################
file_prefix = "histos_g4lbne_"+version+"_"+physics_list+"_"+options.macro+"_"+options.mode;
input_flux_dir = options.input_dir+"/"+version+"/"+physics_list+"/"+options.macro+"/"+options.mode+"/flux/"
output_flux_dir = options.output_dir+"/"+version+"/"+physics_list+"/"+options.macro+"/"+options.mode+"/flux/"

if not os.path.exists(output_flux_dir):
    os.makedirs(output_flux_dir)

plot_files = []
fastmc_files = []
globes_files = []

for file in os.listdir(input_flux_dir+"/histos/"):

    if file.endswith("_"+options.location+".root"):
        plot_files.append(os.path.join(input_flux_dir+"/histos/",file));
    elif file.endswith("_"+options.location+"_fastmc.root"):
        fastmc_files.append(os.path.join(input_flux_dir+"/histos/",file));
    elif file.endswith("_"+options.location+"_globes_flux.txt"):
        globes_files.append(os.path.join(input_flux_dir+"/histos/",file));
    else:
        continue
    
print "Found "+str(len(plot_files))+" histograms to merge."
print "Found "+str(len(fastmc_files))+" fastmc files to merge."
print "Found "+str(len(globes_files))+" globes files to merge."

merge_histos(plot_files,output_flux_dir+"/"+file_prefix+"_"+options.location+".root");
merge_histos(fastmc_files,output_flux_dir+"/"+file_prefix+"_"+options.location+"_fastmc.root");
merge_globes(globes_files,output_flux_dir+"/"+file_prefix+"_"+options.location+"_globes_flux.txt");