add a script to plot signal sums

scripts/plot_sum.py

+#!/bin/python3
+
+import argparse
+import os
+import json
+import ROOT
+import random
+import numpy as np
+from ROOT import gROOT, gInterpreter
+from plot_utils import prepare_canvas, my_style
+
+
+if __name__ == "__main__":
+    # make sure the relative path for root loading is correct
+    owd = os.getcwd()
+    script_dir = os.path.dirname(os.path.realpath(__file__))
+    os.chdir(os.path.join(script_dir, '..'))
+    print('working directory is ' + os.getcwd())
+
+    # get layouts
+    with open(os.path.join(script_dir, '..', 'database', 'channels_layout.json'), 'r') as f:
+        channels_layout = json.load(f)
+    # format keys
+    layout = {key.replace(' ', '_').upper(): channels for key, channels in channels_layout.items()}
+
+    # argument parser
+    parser = argparse.ArgumentParser('draw events')
+    parser.add_argument('root_file', help='a root file of waveform data')
+    parser.add_argument('channels', help='channels to plot, supports {}'.format(list(layout)))
+    parser.add_argument('-o', '--plot-dir', dest='out_dir', default="./plots", help='output directory')
+    parser.add_argument('-t', type=float, default=0.0, dest='thres', help='peak height threshold to sum')
+    args = parser.parse_args()
+
+    if args.channels not in layout:
+        print('Unsupported channel layout {}, please choose one of {}'.format(args.channels, list(layout)))
+        exit(-1)
+
+    # recover paths
+    for attr in ['root_file', 'out_dir']:
+        setattr(args, attr, os.path.join(owd, getattr(args, attr)))
+
+    os.makedirs(args.out_dir, exist_ok=True)
+
+    # batch mode
+    gROOT.SetBatch(True)
+    # root macros
+    gROOT.Macro('rootlogon.C')
+    # declare root graph code
+    func_code = open(os.path.join(script_dir, 'root', 'dataframe_analysis.C')).read()
+    gInterpreter.Declare(func_code)
+
+    # style
+    my_style.cd()
+
+    # root dataframe
+    rdf = ROOT.RDataFrame('EvTree', args.root_file)
+    channels = ['{}_peak'.format(c) for chs in layout[args.channels] for c in chs]
+    nhits = np.arange(1, 10, dtype=int)
+
+
+    rdf = rdf.Define('sum_peak', 'peak_sum({}, {})'.format('{' + ','.join(channels) + '}', args.thres))\
+             .Define('sum_height', 'sum_peak.height')\
+             .Filter('sum_height > 0')
+    spe_est = rdf.Histo1D(ROOT.RDF.TH1DModel('h1', '', 1000, 0, 1000), 'sum_height')
+    spe = spe_est.GetBinCenter(spe_est.GetMaximumBin())
+
+    rdf = rdf.Define('npe', 'sum_height/{}'.format(spe))
+
+    c = ROOT.TCanvas('{} Signal Sum'.format(args.channels.upper()), '', 1920, 1080)
+    pads = prepare_canvas(c, [[nh] for nh in nhits], 'Signal Sums Grouped by N_pixels',
+                          'Number of photo-electrons', 'Counts')
+    hists = []
+    hrange = (45, 0, 45)
+    for nh in nhits:
+        pads[nh].cd()
+        hists.append(
+            rdf.Filter('sum_peak.pos == {}'.format(nh))\
+               .Histo1D(ROOT.RDF.TH1DModel('hsum{}'.format(nh), 'N_pixels = {}'.format(nh), *hrange), 'npe')
+        )
+        hists[-1].SetLineColor(1)
+        hists[-1].SetLineStyle(2)
+        hists[-1].SetFillColorAlpha(int(nh + 1), 0.5)
+        hists[-1].Draw('lf')
+    c.SaveAs(os.path.join(args.out_dir, '{}_sum.png'.format(args.channels.lower())))
+
+    nhits = 2
+    c = ROOT.TCanvas('Total Sum'.format(args.channels.upper()), '', 1920, 1080)
+    pads = prepare_canvas(c, [[nhits]], 'Signal Sums N_pixels >= {}'.format(nhits),
+                          'Number of photo-electrons', 'Counts')
+    pads[nhits].cd()
+    hist = rdf.Filter('sum_peak.pos >= {}'.format(nhits))\
+              .Histo1D(ROOT.RDF.TH1DModel('hsum'.format(nh), '', *hrange), 'npe')
+    hist.SetLineColor(1)
+    hist.SetLineStyle(2)
+    hist.SetFillColorAlpha(38, 0.5)
+    hist.Draw('lf')
+    c.SaveAs(os.path.join(args.out_dir, '{}_sum_ge{}.png'.format(args.channels.lower(), nhits)))
+