Changes to python scripts

benchmarks/clustering/scripts/cluster_plots.py

@@ -50,10 +50,10 @@ def flatten_collection(rdf, collection, cols=None):
 
 def thrown_particles_figure(rdf, save, mcbranch="MCParticles"):
     # define truth particle info
-    dft = flatten_collection(rdf, mcbranch, ['genStatus', 'pdgID', 'ps.x', 'ps.y', 'ps.z', 'mass'])
+    dft = flatten_collection(rdf, mcbranch, ['generatorStatus', 'PDG', 'momentum.x', 'momentum.y', 'momentum.z', 'mass'])
     dft.rename(columns={c: c.replace(mcbranch + '.', '') for c in dft.columns}, inplace=True)
     # select thrown particles
-    dft = dft[dft['genStatus'] == 1]
+    dft = dft[dft['generatorStatus'] == 1]
 
     # figure
     fig, axs = plt.subplots(2, 2, figsize=(16, 12), dpi=120)
@@ -77,7 +77,7 @@ def thrown_particles_figure(rdf, save, mcbranch="MCParticles"):
 
     # calculate kinematics
     get_4vecs = np.vectorize(lambda x, y, z, m: ROOT.Math.PxPyPzMVector(x, y, z, m))
-    fourvecs = get_4vecs(*dft[['ps.x', 'ps.y', 'ps.z', 'mass']].values.T)
+    fourvecs = get_4vecs(*dft[['momentum.x', 'momentum.y', 'momentum.z', 'mass']].values.T)
 
     dft.loc[:, 'p'] = [v.P() for v in fourvecs]
     dft.loc[:, 'eta'] = [v.Eta() for v in fourvecs]

benchmarks/imaging_shower_ML/scripts/check_edep_dists.py

@@ -77,13 +77,13 @@ if __name__ == '__main__':
     rdf = ROOT.RDataFrame("events", args.file)
 
     mc_branch = args.truth_branch
-    dfm = flatten_collection(rdf, mc_branch, ['genStatus', 'pdgID', 'ps.x', 'ps.y', 'ps.z', 'mass'])
+    dfm = flatten_collection(rdf, mc_branch, ['generatorStatus', 'PDG', 'momentum.x', 'momentum.y', 'momentum.z', 'mass'])
     dfm.rename(columns={c: c.replace(mc_branch + '.', '') for c in dfm.columns}, inplace=True)
     # selete incident particles
-    dfm = dfm[dfm['genStatus'].isin([0, 1])]
+    dfm = dfm[dfm['generatorStatus'].isin([0, 1])]
     # NOTE: assumed single particles
     dfm = dfm.groupby('event').first()
-    # p, theta, phi, pT, eta = cartesian_to_polar(*dfm[['ps.x', 'ps.y', 'ps.z']].values.T)
+    # p, theta, phi, pT, eta = cartesian_to_polar(*dfm[['momentum.x', 'momentum.y', 'momentum.z']].values.T)
 
     if args.sim:
         df = flatten_collection(rdf, args.branch, ['energyDeposit'])

benchmarks/imaging_shower_ML/scripts/prepare_tf_dataset.py

@@ -88,10 +88,10 @@ if __name__ == '__main__':
     dfm = flatten_collection(rdf, 'MCParticles', ['generatorStatus', 'PDG', 'momentum.x', 'momentum.y', 'momentum.z', 'mass'])
     dfm.rename(columns={c: c.replace('MCParticles.', '') for c in dfm.columns}, inplace=True)
     # selete incident particles
-    dfm = dfm[dfm['genStatus'].isin([0, 1])]
+    dfm = dfm[dfm['generatorStatus'].isin([0, 1])]
     # NOTE: assumed single particles
     dfm = dfm.groupby('event').first()
-    p, theta, phi, pT, eta = cartesian_to_polar(*dfm[['ps.x', 'ps.y', 'ps.z']].values.T)
+    p, theta, phi, pT, eta = cartesian_to_polar(*dfm[['momentum.x', 'momentum.y', 'momentum.z']].values.T)
     dfm.loc[:, 'p'] = p
     dfm.loc[:, 'theta'] = theta
     dfm.loc[:, 'phi'] = phi

benchmarks/tracking/scripts/tracking_performance.py

@@ -49,10 +49,10 @@ def flatten_collection(rdf, collection, cols=None):
 
 def thrown_particles_figure(rdf, save, mcbranch="MCParticles"):
     # define truth particle info
-    dft = flatten_collection(rdf, mcbranch, ['genStatus', 'pdgID', 'ps.x', 'ps.y', 'ps.z', 'mass'])
+    dft = flatten_collection(rdf, mcbranch, ['generatorStatus', 'PDG', 'momentum.x', 'momentum.y', 'momentum.z', 'mass'])
     dft.rename(columns={c: c.replace(mcbranch + '.', '') for c in dft.columns}, inplace=True)
     # select thrown particles
-    dft = dft[dft['genStatus'] == 1]
+    dft = dft[dft['generatorStatus'] == 1]
 
     # figure
     fig, axs = plt.subplots(2, 2, figsize=(16, 12), dpi=120)
@@ -78,7 +78,7 @@ def thrown_particles_figure(rdf, save, mcbranch="MCParticles"):
 
     # calculate kinematics
     get_4vecs = np.vectorize(lambda x, y, z, m: ROOT.Math.PxPyPzMVector(x, y, z, m))
-    fourvecs = get_4vecs(*dft[['ps.x', 'ps.y', 'ps.z', 'mass']].values.T)
+    fourvecs = get_4vecs(*dft[['momentum.x', 'momentum.y', 'momentum.z', 'mass']].values.T)
 
     dft.loc[:, 'p'] = [v.P() for v in fourvecs]
     dft.loc[:, 'theta'] = [v.Theta() for v in fourvecs]