Coverage for orcasong/plotting/binning.py: 0%

1#!/usr/bin/env python 

2# -*- coding: utf-8 -*- 

3 

4import numpy as np 

5import km3pipe as kp 

6import matplotlib.pyplot as plt 

7 

8import orcasong.modules as modules 

9 

10__author__ = "Stefan Reck" 

11 

12 

13class TimePlotter: 

14 """ 

15 For plotting the time distribution of hits, in a histogram, 

16 and finding a good binning. 

17 

18 The plot will have some bins attached in both directions for 

19 better overview. 

20 

21 Attributes: 

22 ----------- 

23 files : list 

24 The .h5 files to read hits from. 

25 do_mc_hits : bool 

26 Read mchits instead of hits. 

27 det_file : str, optional 

28 Path to a .detx detector geometry file, which can be used to 

29 calibrate the hits. 

30 center_time : bool 

31 Subtract time of first triggered hit from all hit times. Will 

32 also be done for McHits if they are in the blob [default: True]. 

33 add_t0 : bool 

34 If true, add t0 to the time of hits. If using a det_file, 

35 this will already have been done automatically [default: False]. 

36 inactive_du : int, optional 

37 Dont plot hits from this du. 

38 

39 """ 

40 

41 def __init__( 

42 self, 

43 files, 

44 do_mchits=False, 

45 det_file=None, 

46 center_time=True, 

47 add_t0=False, 

48 subtract_t0_mchits=False, 

49 inactive_du=None, 

50 ): 

51 if isinstance(files, str): 

52 self.files = [files] 

53 else: 

54 self.files = files 

55 self.do_mchits = do_mchits 

56 self.det_file = det_file 

57 self.add_t0 = add_t0 

58 self.subtract_t0_mchits = subtract_t0_mchits 

59 self.center_time = center_time 

60 self.inactive_du = inactive_du 

61 

62 self.data = np.array([]) 

63 for file in self.files: 

64 self._read(file, self.det_file) 

65 

66 def hist(self, bins=50, padding=0.5, **kwargs): 

67 """ 

68 Plot the hits as a histogram. 

69 

70 Parameters 

71 ---------- 

72 bins : int or np.array 

73 Number of bins, or bin edges. If bin edges are given, some 

74 bins are attached left and right (next to red vertical lines) 

75 for better overview. 

76 padding : float 

77 Fraction of total number of bins to attach left and right 

78 in the plot. 

79 

80 """ 

81 plt.grid(True, zorder=-10) 

82 plt.xlabel("time") 

83 if self.do_mchits: 

84 plt.ylabel("mchits / bin") 

85 else: 

86 plt.ylabel("hits / bin") 

87 if isinstance(bins, int): 

88 plt.hist(self.data, bins=bins, zorder=10, **kwargs) 

89 else: 

90 self.print_binstats(bins) 

91 plt.hist( 

92 self.data, bins=_get_padded_bins(bins, padding), zorder=10, **kwargs 

93 ) 

94 for bin_line_x in (bins[0], bins[-1]): 

95 plt.axvline(x=bin_line_x, color="firebrick", linestyle="--", zorder=20) 

96 

97 def print_binstats(self, bin_edges): 

98 print(f"Cutoff left: {np.mean(self.data < bin_edges[0]):.2%}") 

99 print(f"Cutoff right: {np.mean(self.data > bin_edges[-1]):.2%}") 

100 print(f"Avg. time per bin: {np.mean(np.diff(bin_edges)):.2f}") 

101 

102 def _read(self, file, det_file=None): 

103 if det_file is not None: 

104 cal = modules.DetApplier(det_file=det_file) 

105 else: 

106 cal = None 

107 if self.center_time or self.add_t0: 

108 time_pp = modules.TimePreproc( 

109 add_t0=self.add_t0, 

110 center_time=self.center_time, 

111 subtract_t0_mchits=self.subtract_t0_mchits, 

112 ) 

113 else: 

114 time_pp = None 

115 

116 pump = kp.io.hdf5.HDF5Pump(filename=file) 

117 for i, blob in enumerate(pump): 

118 if i % 1000 == 0: 

119 print("Blob {}".format(i)) 

120 if cal is not None: 

121 blob = cal.process(blob) 

122 if time_pp is not None: 

123 blob = time_pp.process(blob) 

124 self.data = np.concatenate([self.data, self._get_data_one_event(blob)]) 

125 pump.close() 

126 

127 def _get_data_one_event(self, blob): 

128 if self.do_mchits: 

129 fld = "McHits" 

130 else: 

131 fld = "Hits" 

132 blob_data = blob[fld]["time"] 

133 

134 if self.inactive_du is not None: 

135 dus = blob[fld]["du"] 

136 blob_data = blob_data[dus != self.inactive_du] 

137 return blob_data 

138 

139 

140def _get_padded_bins(bin_edges, padding): 

141 """Add fraction of total # of bins, with same width.""" 

142 bin_width = bin_edges[1] - bin_edges[0] 

143 n_bins = len(bin_edges) - 1 

144 bins_to_add = np.ceil(padding * n_bins) 

145 return np.linspace( 

146 bin_edges[0] - bins_to_add * bin_width, 

147 bin_edges[-1] + bins_to_add * bin_width, 

148 int(n_bins + 2 * bins_to_add + 1), 

149 )