Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2022-11-23 00:00:37

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()