def main():
    # Arguments:
    assert len(sys.argv) > 1
    f_in = sys.argv[1]
    with open("cuts.yaml") as f:
        known_cuts = yaml.load(f)
    cut_key = "pre"
    if len(sys.argv) > 2:
        if sys.argv[2] in known_cuts:
            cut_key = sys.argv[2]
    tcut = TCut(cut_key, known_cuts[cut_key])
    cut = tcut.GetTitle()
    print "[..] Applying the following cut to {}:\n{}\n".format(f_in, cut)
    f_out = f_in.replace(".root", "_{}.root".format(cut_key))

    # Apply the cut:
    rf = root.rfile(f_in)
    tts = rf.get_ttrees()
    tf_out = TFile(f_out, "RECREATE")
    for tt_name, tt in tts.items():
        print "[..] Cutting {}.".format(tt.GetName())
        tt_out = tt.CopyTree(cut)
        tf_out.WriteTObject(tt_out)
    tf_out.Close()
    return True
Example #2
0
def main():
	# Input:
	## Anatuple:
	assert(len(sys.argv) >= 2)
	atuple = root.rfile(sys.argv[1])
	tts = atuple.get_ttrees()
	
	## Plots to make:
	with open("plotter_info.yaml", 'r') as f:
		plot_info = yaml.load(f)
	
	## Cuts to make:
	cuts_dict = get_cuts()
	
	# Output:
	if len(sys.argv) >= 3:
		f_out = sys.argv[2]
	else:
		f_out = atuple.name.replace("anatuple", "plots")
	tf_out = TFile(f_out, "RECREATE")
	
	# Start:
	gROOT.SetBatch()
	
	# Write cuts:
	## (I do this seperately so they appear first in the file list.)
	for key, tcut in cuts_dict.items():
		tcut.Write()
	
	# Write histograms:
	for cut_key, tcut in cuts_dict.items():
		print "[..] Making plots for {}, which is\n\t{}".format(cut_key, tcut.GetTitle())
		# TH1s:
		for var_name, info in plot_info["th1"].items():
			for tt in tts:
				tt_name = tt.GetName()
				print "\t[..] Plotting {} for {}.".format(info["var"], tt_name)
				th_name = "{}_{}_{}".format(tt_name, var_name, cut_key)
#				print tt_name, var
				tt.Draw("{}>>{}({}, {}, {})".format(info["var"], th_name, info["binning"][0], info["binning"][1], info["binning"][2]), tcut)
				th = gDirectory.Get(th_name)
				th.Write()
		# TH2s:
		for var_name, info in plot_info["th2"].items():
			for tt in tts:
				tt_name = tt.GetName()
				print "\t[..] Plotting {} for {}.".format(info["var"], tt_name)
				th_name = "{}_{}_cut{}".format(tt_name, var_name, cut_key)
#				print tt_name, var
				tt.Draw("{}:{}>>{}({}, {}, {}, {}, {}, {})".format(info["var"][0], info["var"][1], th_name, info["binning"][0], info["binning"][1], info["binning"][2], info["binning"][3], info["binning"][4], info["binning"][5]), tcut)
				th = gDirectory.Get(th_name)
				th.Write()
	print "[OK] Plots sucessfully written to\n\t{}".format(f_out)
	return True
def main():
	# Arguments:
	assert len(sys.argv) > 1
	f_in = sys.argv[1]
	known_cuts = get_cuts()
	cut_key = "preselAndMasy"
	if len(sys.argv) > 2:
		if sys.argv[2] in known_cuts:
			cut_key = sys.argv[2]
	cut = known_cuts[cut_key].GetTitle()
	f_out = f_in.replace(".root", "_{}.root".format(cut_key))
	
	# Apply the cut:
	rf = root.rfile(f_in)
	tts = rf.get_ttrees()
	tf_out = TFile(f_out, "RECREATE")
	for tt in tts:
		print "Cutting {} ...".format(tt.GetName())
		tt_out = tt.CopyTree(cut)
		tf_out.WriteTObject(tt_out)
	tf_out.Close()
	return True