Lb2Lcpi_pidgen.py 4.68 KB
Newer Older
1
import os
2
3
4
5
6
import sys

## START OF CONFIG
# Read comments and check vars
# at least until end of config section
7

Konstantin Gizdov's avatar
Konstantin Gizdov committed
8
# List of input ROOT files with MC ntuples. Format:
9
10
#   (inputfile, outputfile, dataset)
files = [
Konstantin Gizdov's avatar
Konstantin Gizdov committed
11
  ("root://eoslhcb.cern.ch//eos/lhcb/wg/PID/PIDGen/Validation/Lb2Lcpi/sim08_2012_md.root", "sim08_2012_md_pidgen.root", "MagDown_2012"),
12
13
14
15
16
17
18
19
20
21
]

# Name of the input tree
# Could also include ROOT directory, e.g. "Dir/Ntuple"
input_tree = "lb2lch_filt"

# Postfixes of the Pt, Eta and Ntracks variables (ntuple variable name w/o branch name)
# e.g. if the ntuple contains "pion_PT", it should be just "PT"
ptvar  = "pt"
etavar = "eta"
22
23
24
25
26
pvar   = None
## Could use P variable instead of eta
# etavar = None
# pvar   = "p"

Konstantin Gizdov's avatar
Konstantin Gizdov committed
27
ntrvar = "nTracks"  # This should correspond to the number of "Best tracks", not "Long tracks"!
28

29
seed = None   # No initial seed
30
# seed = 1    # Alternatively, could set initial random seed
31

32
# Dictionary of tracks with their PID variables, in the form {branch name}:{pidvars}
Konstantin Gizdov's avatar
Konstantin Gizdov committed
33
34
35
36
37
# For each track branch name, {pidvars} is a dictionary in the form {ntuple variable}:{pid config},
#   where
#     {ntuple variable} is the name of the corresponding ntuple PID variable without branch name,
#   and
#     {pid_config} is the string describing the PID configuration.
38
39
40
# Run PIDCorr.py without arguments to get the full list of PID configs
tracks = {
  'h'   : {
Konstantin Gizdov's avatar
Konstantin Gizdov committed
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
            "V2ProbNNK"  : "pi_V2ProbNNK",
            "V2ProbNNpi" : "pi_V2ProbNNpi",
            "V2ProbNNp"  : "pi_V2ProbNNp",
            "V3ProbNNK"  : "pi_V3ProbNNK",
            "V3ProbNNpi" : "pi_V3ProbNNpi",
            "V3ProbNNp"  : "pi_V3ProbNNp",
            "pidk"       : "pi_CombDLLK",
            "pidp"       : "pi_CombDLLp",
           },
  'lpi' : {
            "V2ProbNNK"  : "pi_V2ProbNNK",
            "V2ProbNNpi" : "pi_V2ProbNNpi",
            "V2ProbNNp"  : "pi_V2ProbNNp",
            "V3ProbNNK"  : "pi_V3ProbNNK",
            "V3ProbNNpi" : "pi_V3ProbNNpi",
            "V3ProbNNp"  : "pi_V3ProbNNp",
            "pidk"       : "pi_CombDLLK",
            "pidp"       : "pi_CombDLLp",
           },
  'lk'  : {
            "V2ProbNNK"  : "K_V2ProbNNK",
            "V2ProbNNpi" : "K_V2ProbNNpi",
            "V2ProbNNp"  : "K_V2ProbNNp",
            "V3ProbNNK"  : "K_V3ProbNNK",
            "V3ProbNNpi" : "K_V3ProbNNpi",
            "V3ProbNNp"  : "K_V3ProbNNp",
            "pidk"       : "K_CombDLLK",
            "pidp"       : "K_CombDLLp",
           },
70
  'lp'  : {
Konstantin Gizdov's avatar
Konstantin Gizdov committed
71
72
73
74
75
76
77
78
79
            "V2ProbNNK"  : "p_V2ProbNNK",
            "V2ProbNNpi" : "p_V2ProbNNpi",
            "V2ProbNNp"  : "p_V2ProbNNp",
            "V3ProbNNK"  : "p_V3ProbNNK",
            "V3ProbNNpi" : "p_V3ProbNNpi",
            "V3ProbNNp"  : "p_V3ProbNNp",
            "pidk"       : "p_CombDLLK",
            "pidp"       : "p_CombDLLp",
           },
80
81
}

82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
# IF ON LXPLUS: if /tmp exists and is accessible, use for faster processing
# IF NOT: use /tmp if you have enough RAM
# temp_folder = '/tmp'
# ELSE: use current folder
temp_folder = '.'

## END OF CONFIG


# make sure we don't overwrite local files and prefix them with random strings
import string
import random
rand_string = ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(10))  # get 10 random chars for temp_file prefix

temp_file_prefix = temp_folder + '/' + rand_string  # prefix temp files with folder and unique ID

98
99
output_tree = input_tree.split("/")[-1]

Konstantin Gizdov's avatar
Konstantin Gizdov committed
100
for input_file, output_file, dataset in files :
101
  tmpinfile = input_file
102
  tmpoutfile = "%s_tmp1.root" % temp_file_prefix
Konstantin Gizdov's avatar
Konstantin Gizdov committed
103
104
  for track, subst in tracks.iteritems() :
    for var, config in subst.iteritems() :
105
106
      command = "python $PIDPERFSCRIPTSROOT/scripts/python/PIDGenUser/PIDGen.py"
      command += " -m %s_%s" % (track, ptvar)
107
      if etavar:
108
        command += " -e %s_%s" % (track, etavar)
109
110
111
112
113
      elif pvar:
        command += " -q %s_%s" % (track, pvar)
      else:
        print('Specify either ETA or P branch name per track')
        sys.exit(1)
114
      command += " -n %s" % ntrvar
115
      command += " -t %s" % input_tree
116
117
118
119
120
      command += " -p %s_%s_corr" % (track, var)
      command += " -c %s" % config
      command += " -d %s" % dataset
      command += " -i %s" % tmpinfile
      command += " -o %s" % tmpoutfile
Konstantin Gizdov's avatar
Konstantin Gizdov committed
121
      if seed :
122
        command += " -s %d" % seed
123
124

      tmpinfile = tmpoutfile
125
126
      if 'tmp1' in tmpoutfile:
        tmpoutfile = tmpoutfile.replace('tmp1', 'tmp2')
Konstantin Gizdov's avatar
Konstantin Gizdov committed
127
      else :
128
        tmpoutfile = tmpoutfile.replace('tmp2', 'tmp1')
129

130
      print(command)
131
132
      os.system(command)

133
134
135
136
137
  if "root://" in output_file:
    print("xrdcp %s %s" % (tmpinfile, output_file))
    os.system("xrdcp %s %s" % (tmpinfile, output_file))
  else:
    print("mv %s %s" % (tmpinfile, output_file))
Konstantin Gizdov's avatar
Konstantin Gizdov committed
138
    os.system("mv %s %s" % (tmpinfile, output_file))