Skip to content
Snippets Groups Projects
rich_samples.py 1.36 KiB
Newer Older
  • Learn to ignore specific revisions
  • import ROOT
    import numpy as np
    from ROOT import gROOT
    import argparse
    from matplotlib import pyplot as plt
    from matplotlib.patches import Ellipse
    import matplotlib.transforms as transforms
    
    parser = argparse.ArgumentParser(description='rich reconstruction analysis')
    parser.add_argument('nevents', type=int, help='number of events')
    args = parser.parse_args()
    
    gROOT.Macro('rootlogon.C')
    f = ROOT.TFile("rich_test_reco.root")
    events = f.events
    
    nev = 0
    while nev < args.nevents:
        iev = np.random.randint(0, events.GetEntries())
        events.GetEntry(iev)
        if events.RICHClusters.size() == 0:
            continue
    
        nev += 1
        x, y = [], []
        for hit in events.RecoForwardRICHHits:
            x.append(hit.position.x)
            y.append(hit.position.y)
    
        fig, ax = plt.subplots(figsize=(9, 9), dpi=120)
        colors = plt.rcParams['axes.prop_cycle'].by_key()['color']
        ax.scatter(x, y, s=9, marker='s', color=colors[0])
    
        for cluster in events.RICHClusters:
            ellipse = Ellipse((cluster.position.x, cluster.position.y), width=cluster.radius * 2, height=cluster.radius * 2,
                              edgecolor='red', fill=False, lw=2, alpha=0.8)
            ellipse.set_transform(ax.transData)
            ax.add_patch(ellipse)
        ax.set_xlim(-60, 60)
        ax.set_ylim(-60, 60)
        ax.set_xlabel('X (mm)')
        ax.set_ylabel('Y (mm)')
        fig.savefig('rich_cluster_{}.png'.format(iev))