safeUncalculable/tools.py at main · sambt/safeUncalculable · 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
import numpy as np
import matplotlib.pyplot as plt
import time
import h5py
from sklearn.metrics import roc_auc_score, roc_curve
from multiprocessing import Process

def get_data(file):
    with h5py.File(file,"r") as f:
        z = f['jet1_constit_z'][()]
        eta = f['jet1_constit_eta'][()]
        phi = f['jet1_constit_phi'][()]

        pz = f['pjet1_constit_z'][()]
        peta = f['pjet1_constit_eta'][()]
        pphi = f['pjet1_constit_phi'][()]

    X1 = np.concatenate((z[:,:,np.newaxis],eta[:,:,np.newaxis],phi[:,:,np.newaxis]),axis=-1)
    Y1 = np.ones(z.shape[0])

    X2 = np.concatenate((pz[:,:,np.newaxis],peta[:,:,np.newaxis],pphi[:,:,np.newaxis]),axis=-1)
    Y2 = np.zeros(pz.shape[0])

    X = np.concatenate((X1,X2),axis=0)
    Y = np.concatenate((Y1,Y2),axis=0)

    del X1,Y1,X2,Y2

    return X,Y

def get_data_softDrop(file):
    with h5py.File(file,"r") as f:
        z = f['jet1_sd_constit_z'][()]
        eta = f['jet1_sd_constit_eta'][()]
        phi = f['jet1_sd_constit_phi'][()]

        pz = f['pjet1_sd_constit_z'][()]
        peta = f['pjet1_sd_constit_eta'][()]
        pphi = f['pjet1_sd_constit_phi'][()]

    X1 = np.concatenate((z[:,:,np.newaxis],eta[:,:,np.newaxis],phi[:,:,np.newaxis]),axis=-1)
    Y1 = np.ones(z.shape[0])

    X2 = np.concatenate((pz[:,:,np.newaxis],peta[:,:,np.newaxis],pphi[:,:,np.newaxis]),axis=-1)
    Y2 = np.zeros(pz.shape[0])

    X = np.concatenate((X1,X2),axis=0)
    Y = np.concatenate((Y1,Y2),axis=0)

    del X1,Y1,X2,Y2

    return X,Y