update draw_clusters

benchmarks/sampling_ecal/scripts/draw_cluster.py

@@ -132,21 +132,38 @@ if __name__ == '__main__':
     # read data
     load_root_macros(args.macros)
     df = get_hits_data(args.file, args.iev, branch=args.branch)
-    dfmcp = get_mcp_simple(args.file, args.iev, 'mcparticles2')
-    vec = dfmcp[['px', 'py', 'pz']].values[0]
-    vec = vec/np.linalg.norm(vec)
-    # df = df[df['cluster'] == args.icl]
+    df = df[df['cluster'] == args.icl]
     if not len(df):
         print("Error: do not find any hits for cluster {:d} in event {:d}".format(args.icl, args.iev))
         exit(-1)
+    # convert to polar coordinates (mrad), and stack all r values
+    df['r'] = np.sqrt(df['x'].values**2 + df['y'].values**2 + df['z'].values**2)
+    df['phi'] = np.arctan2(df['y'].values, df['x'].values)*1000.
+    df['theta'] = np.arccos(df['z'].values/df['r'].values)*1000.
+    df['eta'] = -np.log(np.tan(df['theta'].values/1000./2.))
+
+    # truth
+    dfmcp = get_mcp_simple(args.file, args.iev, 'mcparticles2').iloc[0]
+    pdgbase = ROOT.TDatabasePDG()
+    inpart = pdgbase.GetParticle(int(dfmcp['pid']))
+    print("Incoming particle = {}, pdgcode = {}, charge = {}, mass = {}"\
+          .format(inpart.GetName(), inpart.PdgCode(), inpart.Charge(), inpart.Mass()))
+    # neutral particle, no need to consider magnetic field
+    if np.isclose(inpart.Charge(), 0., rtol=1e-5):
+        vec = dfmcp[['px', 'py', 'pz']].values
+    # charge particle, use the cluster center
+    else:
+        flayer = df[df['layer'] == df['layer'].min()]
+        vec = flayer[['x', 'y', 'z']].mean().values
+    vec = vec/np.linalg.norm(vec)
+
     # particle line from (0, 0, 0) to the inner Ecal surface
     length = rmin/np.sqrt(vec[0]**2 + vec[1]**2)
     pline = np.transpose(vec*np.mgrid[0:length:2j][:, np.newaxis])
-
+    cmap = truncate_colormap(plt.get_cmap('jet'), 0.1, 0.9)
 
     os.makedirs(args.outdir, exist_ok=True)
     # cluster plot
-    cmap = truncate_colormap(plt.get_cmap('jet'), 0.1, 0.9)
     fig = plt.figure(figsize=(15, 12), dpi=160)
     ax = fig.add_subplot(111, projection='3d')
     # draw particle line
@@ -183,12 +200,6 @@ if __name__ == '__main__':
 
 
     # projection plot
-    # convert to polar coordinates (mrad), and stack all r values
-    df['r'] = np.sqrt(df['x'].values**2 + df['y'].values**2 + df['z'].values**2)
-    df['phi'] = np.arctan2(df['y'].values, df['x'].values)*1000.
-    df['theta'] = np.arccos(df['z'].values/df['r'].values)*1000.
-    df['eta'] = -np.log(np.tan(df['theta'].values/1000./2.))
-
     # convert to mrad
     vecp = np.asarray([np.arccos(vec[2]), np.arctan2(vec[1], vec[0])])*1000.
     phi_rg = np.asarray([vecp[1] - args.topo_range, vecp[1] + args.topo_range])

benchmarks/sampling_ecal/scripts/draw_cluster_layers.py

@@ -87,15 +87,21 @@ if __name__ == '__main__':
     load_root_macros(args.macros)
     df = get_hits_data(args.file, args.iev, args.branch)
     df = df[df['cluster'] == args.icl]
+    if not len(df):
+        print("Error: do not find any hits for cluster {:d} in event {:d}".format(args.icl, args.iev))
+        exit(-1)
     # convert to polar coordinates (mrad), and stack all r values
     df['r'] = np.sqrt(df['x'].values**2 + df['y'].values**2 + df['z'].values**2)
     df['phi'] = np.arctan2(df['y'].values, df['x'].values)*1000.
     df['theta'] = np.arccos(df['z'].values/df['r'].values)*1000.
     df['eta'] = -np.log(np.tan(df['theta'].values/1000./2.))
+
+    # truth
     dfmcp = get_mcp_simple(args.file, args.iev, 'mcparticles2').iloc[0]
     pdgbase = ROOT.TDatabasePDG()
     inpart = pdgbase.GetParticle(int(dfmcp['pid']))
-    print("Incoming particle = {}, pdgcode = {}, charge = {}, mass = {}".format(inpart.GetName(), inpart.PdgCode(), inpart.Charge(), inpart.Mass()))
+    print("Incoming particle = {}, pdgcode = {}, charge = {}, mass = {}"\
+          .format(inpart.GetName(), inpart.PdgCode(), inpart.Charge(), inpart.Mass()))
     # neutral particle, no need to consider magnetic field
     if np.isclose(inpart.Charge(), 0., rtol=1e-5):
         vec = dfmcp[['px', 'py', 'pz']].values
@@ -105,15 +111,11 @@ if __name__ == '__main__':
         vec = flayer[['x', 'y', 'z']].mean().values
     vec = vec/np.linalg.norm(vec)
 
-    if not len(df):
-        print("Error: do not find any hits for cluster {:d} in event {:d}".format(args.icl, args.iev))
-        exit(-1)
     # particle line from (0, 0, 0) to the inner Ecal surface
     length = rmin/np.sqrt(vec[0]**2 + vec[1]**2)
     pline = np.transpose(vec*np.mgrid[0:length:2j][:, np.newaxis])
     cmap = truncate_colormap(plt.get_cmap('jet'), 0.1, 0.9)
 
-
     # convert truth to mrad
     vecp = np.asarray([np.arccos(vec[2]), np.arctan2(vec[1], vec[0])])*1000.
     phi_rg = np.asarray([vecp[1] - args.topo_range, vecp[1] + args.topo_range])