GEMCSCSegment/test/plotGEMCSCCoincidenceRate.py

0001 #!/usr/bin/env python3.9
0002 from __future__ import annotations
0003 from dataclasses import dataclass
0004 from functools import cached_property
0005 from typing import Optional
0006 import argparse
0007 from pathlib import Path
0008 import numpy as np
0009 import uproot
0010 import hist
0011 from hist import intervals
0012 import awkward as ak
0013 import matplotlib as mpl
0014 import matplotlib.pyplot as plt
0015 import mplhep as hep
0016 hep.style.use(hep.style.CMS)
0017
0018 __author__ = "Seungjin Yang"
0019 __email__ = "seungjin.yang@cern.ch"
0020 __version__ = "CMSSW_12_6_0_pre2"
0021
0022 @dataclass
0023 class Efficiency1D:
0024     x: np.ndarray
0025     y: np.ndarray
0026     ylow: np.ndarray
0027     yup: np.ndarray
0028
0029     @classmethod
0030     def from_hist(cls, h_num, h_den) -> Efficiency1D:
0031         num = h_num.values()
0032         den = h_den.values()
0033
0034         x = h_den.axes[0].centers
0035         y = np.divide(num, den, where=den > 0)
0036         ylow, yup = intervals.clopper_pearson_interval(num, den)
0037         return cls(x, y, ylow, yup)
0038
0039     @cached_property
0040     def yerr_low(self) -> np.ndarray:
0041         return self.y - self.ylow
0042
0043     @cached_property
0044     def yerr_up(self) -> np.ndarray:
0045         return self.yup - self.y
0046
0047     @property
0048     def yerr(self) -> tuple[np.ndarray, np.ndarray]:
0049         return (self.yerr_low, self.yerr_up)
0050
0051     def plot(self, ax: Optional[mpl.axes.Subplot] = None, **kwargs):
0052         ax = ax or plt.gca()
0053         return ax.errorbar(self.x, self.y, self.yerr, **kwargs)
0054
0055 def plot_eff(eff_twofold_all: Efficiency1D,
0056         eff_twofold_muon: Efficiency1D,
0057         eff_threefold_all: Efficiency1D,
0058         eff_threefold_muon: Efficiency1D,
0059         gem_lable: str,
0060         output_dir: Path):
0061     fig, ax = plt.subplots(figsize=(16, 8))
0062
0063     eff_twofold_all.plot(ax=ax, label='Twofold', ls='', marker='s')
0064     eff_twofold_muon.plot(ax=ax, label='Twofold & Muon', ls='', marker='s')
0065     eff_threefold_all.plot(ax=ax, label='Threefold', ls='', marker='o')
0066     eff_threefold_muon.plot(ax=ax, label='Threefold & Muon', ls='', marker='o')
0067
0068     ax.set_xlabel('Chamber')
0069     ax.set_ylabel('Efficiency')
0070
0071     ax.set_xticks(eff_twofold_all.x)
0072     ax.grid()
0073     ax.set_ylim(0.5, 1)
0074     ax.legend(title=gem_lable, ncol=2, loc='lower center')
0075
0076     fig.tight_layout()
0077
0078     output_path = output_dir / gem_lable
0079     for suffix in ['.png', '.pdf']:
0080         fig.savefig(output_path.with_suffix(suffix))
0081
0082 def name_gem_layer(region: int, station: int, layer: int) -> str:
0083     if region == 1:
0084         region_char = 'P'
0085     elif region == -1:
0086         region_char = 'M'
0087     else:
0088         raise ValueError(f'{region=:}')
0089     return f'GE{station}1-{region_char}-L{layer}'
0090
0091
0092 def plot_layer(tree, region: int, station: int, output_dir: Path):
0093     expressions = [
0094         'region',
0095         'station',
0096         'chamber',
0097         'csc_is_muon',
0098         'gem_layer'
0099     ]
0100
0101     cut_list = [
0102         f'region == {region}',
0103         f'station == {station}',
0104         'ring == 1',
0105         'csc_num_hit == 6',
0106         '~gem_chamber_has_error',
0107     ]
0108     cut = ' & '.join(f'({each})' for each in cut_list)
0109
0110     num_chambers = 36 if station == 1 else 18
0111     chamber_axis = hist.axis.Regular(num_chambers, 0.5, num_chambers + 0.5, name='chamber', label='Chamber')
0112     muon_axis = hist.axis.hist.axis.Boolean(name='muon')
0113
0114     # has a csc segment
0115     h_csc = hist.Hist(chamber_axis, muon_axis, storage=hist.storage.Int64())
0116     # has a CSC segment and a hit on L1(L2) layer
0117     h_l1 = h_csc.copy()
0118     h_l2 = h_csc.copy()
0119     # has a csc segment and hits on both layers
0120     h_both = h_csc.copy()
0121
0122     for chunk in tree.iterate(expressions, cut=cut):
0123         has_l1 = ak.any(chunk.gem_layer == 1, axis=1)
0124         has_l2 = ak.any(chunk.gem_layer == 2, axis=1)
0125         has_both = has_l1 & has_l2
0126
0127         h_csc.fill(chamber=chunk.chamber, muon=chunk.csc_is_muon)
0128         h_l1.fill(chamber=chunk.chamber[has_l1], muon=chunk.csc_is_muon[has_l1])
0129         h_l2.fill(chamber=chunk.chamber[has_l2], muon=chunk.csc_is_muon[has_l2])
0130         h_both.fill(chamber=chunk.chamber[has_both], muon=chunk.csc_is_muon[has_both])
0131
0132     # twofold efficiencies
0133     eff_l1_twofold_all = Efficiency1D.from_hist(h_num=h_l1.project(0), h_den=h_csc.project(0))
0134     eff_l2_twofold_all = Efficiency1D.from_hist(h_num=h_l2.project(0), h_den=h_csc.project(0))
0135
0136     # threefold efficiencies
0137     eff_l1_threefold_all = Efficiency1D.from_hist(h_num=h_both.project(0), h_den=h_l2.project(0))
0138     eff_l2_threefold_all = Efficiency1D.from_hist(h_num=h_both.project(0), h_den=h_l1.project(0))
0139
0140     # efficiencies
0141     muon_mask = (slice(None), True)
0142     eff_l1_twofold_muon = Efficiency1D.from_hist(h_num=h_l1[muon_mask], h_den=h_csc[muon_mask])
0143     eff_l2_twofold_muon = Efficiency1D.from_hist(h_num=h_l2[muon_mask], h_den=h_csc[muon_mask])
0144     eff_l1_threefold_muon = Efficiency1D.from_hist(h_num=h_both[muon_mask], h_den=h_l2[muon_mask])
0145     eff_l2_threefold_muon = Efficiency1D.from_hist(h_num=h_both[muon_mask], h_den=h_l1[muon_mask])
0146
0147     plot_eff(eff_twofold_all=eff_l1_twofold_all,
0148              eff_twofold_muon=eff_l1_twofold_muon,
0149              eff_threefold_all=eff_l1_threefold_all,
0150              eff_threefold_muon=eff_l1_threefold_muon,
0151              gem_lable=name_gem_layer(region, station, layer=1),
0152              output_dir=output_dir)
0153
0154     plot_eff(eff_twofold_all=eff_l2_twofold_all,
0155              eff_twofold_muon=eff_l2_twofold_muon,
0156              eff_threefold_all=eff_l2_threefold_all,
0157              eff_threefold_muon=eff_l2_threefold_muon,
0158              gem_lable=name_gem_layer(region, station, layer=2),
0159              output_dir=output_dir)
0160
0161
0162 def main():
0163     parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
0164     parser.add_argument('-i', '--input-path', type=Path, required=True, help='input file')
0165     parser.add_argument('-t', '--treepath', type=str, default='gemcscCoincidenceRateAnalyzer/gemcsc', help='tree path')
0166     parser.add_argument('-o', '--output-dir', type=Path, default=Path.cwd(), help='output directory')
0167     parser.add_argument('--ge11', action=argparse.BooleanOptionalAction, default=True, help='plot GE11')
0168     parser.add_argument('--ge21', action=argparse.BooleanOptionalAction, default=True, help='plot GE21')
0169     args = parser.parse_args()
0170
0171     if not args.ge11 and not args.ge21:
0172         raise RuntimeError
0173
0174     tree = uproot.open(f'{args.input_path}:{args.treepath}')
0175
0176     region_list = [-1, 1]
0177     station_list = []
0178     if args.ge11:
0179         station_list.append(1)
0180     if args.ge21:
0181         station_list.append(2)
0182
0183     if not args.output_dir.exists():
0184         print(f'mkdir -p {args.output_dir}')
0185         args.output_dir.mkdir(parents=True)
0186
0187     for region in region_list:
0188         for station in station_list:
0189             print(f'plotting layers in {region=} & {station=}')
0190             plot_layer(tree=tree, region=region, station=station, output_dir=args.output_dir)
0191
0192 if __name__ == "__main__":
0193     main()