Combine-Template-Creation/MakeInputRoot_OnShell.py at main · hexutils/Combine-Template-Creation · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
#Code written by Jeffrey Davis, JHU Graduate Student

import os, glob
import sys
import ROOT

# What this script does is trim and rename the template files and prepare them for the datacards.
# The expected input is a list of templates that would belong to the sample datacard.

#========================== How to run this ===========================
# 1. First argument is the name of the template you want to create

# 2. The next arguments are all of the input root files with the histrograms needed
def Make_Template_With_Fake_Data(OutName,names):
  # OutName = Output Root File Name
  # names = Input Root Files

  hists = [] # Holds the histograms to add to the final TFile
  used_names = [] # Stores the names of the histrograms

  for nm in names:
    fin = ROOT.TFile.Open(nm)
    for key in fin.GetListOfKeys():
      if "TH1" in key.GetClassName():
        h_name = key.GetName()
        h_temp = fin.Get(h_name)
        h_temp.SetDirectory(0)
        if h_name not in used_names:
          print(h_name)
          hists.append(h_temp)
          used_names.append(h_name)
    fin.Close()

  # Make a fake data histogram #

  Fake_Data = []
  for hist in hists:
    h_name = hist.GetName()
    if "0PM" in h_name:
      Fake_Data.append(hist)
    if "bkg" in h_name:
      Fake_Data.append(hist)

  Fake_Data_Hist = Fake_Data[0].Clone("data_obs")
  for i in range(1,len(Fake_Data)):
    Fake_Data_Hist.Add(Fake_Data[i])

  fout =  ROOT.TFile.Open(OutName,"recreate")
  fout.cd()

  Fake_Data_Hist.Write("data_obs")

  for hist in hists:
    if ("bkg_ew_negative" not in hist.GetName()):
      if ("bkg_ew_positive" in hist.GetName()):
        hist.Write("bkg_ew")
      else:
        hist.Write()

  print(OutName)

def main():
  output_dir = sys.argv[2]
  if not os.path.exists(output_dir):
      os.mkdir(output_dir)
  if output_dir[0] == "/":
      output_dir = output_dir.strip("/")
      output_dir = "/" + output_dir
  else:
      output_dir = output_dir.strip("/")

  Input_Dir = sys.argv[1]
  for filename in glob.iglob(Input_Dir+'/**', recursive=True):
    if os.path.isfile(filename) and (".root" in filename):
      out_ext=filename
      if "/" in filename:
        out_ext = filename.split("/")[-1]
        out_ext = out_ext.split(".")[0]+".input."+out_ext.split(".")[1]
      Make_Template_With_Fake_Data(output_dir+"/"+out_ext,[filename])

if __name__ == "__main__":
    main()