forked from tkopfMaster/Transfer_NN_Recoil
-
Notifications
You must be signed in to change notification settings - Fork 0
/
prepareOutputFile.py
executable file
·57 lines (46 loc) · 2.3 KB
/
prepareOutputFile.py
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
#!/usr/bin/env python
import h5py
import sys
import root_numpy as rnp
import pandas as pd
import numpy as np
from os import environ
from getNNinputs import kar2pol, pol2kar, angularrange
from Training import loadInputsTargetsWeights
import matplotlib as mpl
import matplotlib.pyplot as plt
def prepareOutput(outputD, ll, plotsD, rootFile, NN_MVA):
NN_Output = h5py.File("%sNN_Output_applied_%s.h5"%(outputD,ll), "r+")
mZ_x, mZ_y = (NN_Output['MET_GroundTruth'][:,0]), (NN_Output['MET_GroundTruth'][:,1])
a_x, a_y = (NN_Output['MET_Predictions'][:,0]), (NN_Output['MET_Predictions'][:,1])
mZ_r, mZ_phi = kar2pol(mZ_x, mZ_y)
mZ_r = NN_Output['Boson_Pt'][:]
a_r, a_phi = kar2pol(a_x, a_y)
NN_LongZ, NN_PerpZ = -np.cos(angularrange(np.add(a_phi,-mZ_phi)))*a_r, np.sin(angularrange(a_phi-mZ_phi))*a_r
#HDF5
dset = NN_MVA.create_dataset("NN_LongZ", dtype='d', data=NN_LongZ)
dset1 = NN_MVA.create_dataset("NN_PerpZ", dtype='d', data=NN_PerpZ)
dset2 = NN_MVA.create_dataset("NN_Phi", dtype='d', data=a_phi)
dset3 = NN_MVA.create_dataset("NN_Pt", dtype='d', data=a_r)
dset4 = NN_MVA.create_dataset("Boson_Pt", dtype='d', data=mZ_r)
dset5 = NN_MVA.create_dataset("NN_x", dtype='d', data=a_x)
dset6 = NN_MVA.create_dataset("NN_y", dtype='d', data=a_y)
dset7 = NN_MVA.create_dataset("Boson_x", dtype='d', data=mZ_x)
dset8 = NN_MVA.create_dataset("Boson_y", dtype='d', data=mZ_y)
dset9 = NN_MVA.create_dataset("Boson_Phi", dtype='d', data=mZ_phi)
NN_MVA.close()
#Root
#treename = ll+"_nominal/ntuple"
#Root_array = rnp.root2array(rootFile, treename=treename)
#print("shape Root_array", Root_array.shape)
NN_array = np.array([(a_r[i], a_phi[i], a_x[i], a_y[i], NN_LongZ[i], NN_PerpZ[i]) for i in range(len(a_r))],
dtype=[('NN_Pt', np.float32), ('NN_Phi', np.float32), ('NN_x', np.float32), ('NN_y', np.float32), ('NN_LongZ', np.float32), ('NN_PerpZ', np.float32)])
print("shape NN_array", NN_array.shape)
rnp.array2root(NN_array, "%s/NN_MVA_%s.root"%(outputDir,ll), mode='recreate')
if __name__ == "__main__":
outputDir = sys.argv[1]
ll = sys.argv[2]
plotsD = sys.argv[3]
rootfile = sys.argv[4]
NN_MVAfile = h5py.File("%s/NN_MVA_%s.h5"%(outputDir,ll), "w")
prepareOutput(outputDir, ll, plotsD, rootfile, NN_MVAfile)