36 def __init__(self, path: str, key: str, particlelist: list, labels: list):
38 Initialize the dataloader for PID prior training.
41 path (str): Path to the root file containing the data.
42 key (str): Key (i.e. path) of the tree within the root file.
43 particlelist (list(int)): List of particle PDG values for which the model has to be trained.
44 labels (str): Labels of pandas columns containing cos(theta), momentum and PDG values (in this order).
48 data = data[key].pandas.df(labels)
49 df = data.dropna().reset_index(drop=
True)
50 df.loc[:, labels[2]] = df.loc[:, labels[2]].abs()
51 droplist = np.setdiff1d(np.unique(df[labels[2]].values), particlelist)
53 df = df.drop(df.loc[df[labels[2]] == i].index).reset_index(drop=
True)
55 x = np.hstack((x, (np.sin(np.arccos(x[:, 0])) * x[:, 1]).reshape(-1, 1)))
56 pol = PolynomialFeatures(2, include_bias=
False)
57 x = pol.fit_transform(x)
59 self.
x = x.astype(
"float32")
62 y = le.fit_transform(y)
64 self.
y = y.astype(
"int64")