Coverage for src/pyhiperta/R0Reader.py: 99%

1# Copyright 2026 CNRS 

2# This software is distributed under the terms of the CeCILL-C free software license. 

3 

4"""Implements the reading/loading of R0 data (waveforms) and service data (gains/pedestals, etc.) from R0 hdf5 files.""" 

5 

6from collections.abc import Iterable 

7from pathlib import Path 

8 

9import numpy as np 

10import tables 

11 

12 

13class R0HDF5Dataset: 

14 """HDF5 reader class inspired by pytorch datasets. 

15 

16 The main purpose of this class is to load shower waveforms in batch from hdf5 files, allowing to 

17 seamlessly iterate over batches independently of the number of files or the number of events per file. 

18 

19 Parameters 

20 ---------- 

21 R0_files : Path or str or Iterable[Path] or Iterable[str] 

22 If `R0_files` is a string or a Path, it is expected to be the path to a single ".h5" R0 file, or to 

23 be the path to a folder containing the ".h5" R0 files to use. 

24 If it is an Iterable, each element is expected to be a string or path to a ".h5" data file. 

25 batch_size : int 

26 Number of shower events to stack to create a batch. 

27 nb_first_slice_to_reject : int, optionnal 

28 If provided, the first `nb_first_slice_to_reject` frames of the shower videos will be discarded when 

29 loading the waveforms. (This is required at the moment for LST data) 

30 nb_last_slice_to_reject : int, optionnal 

31 If provided, the last `nb_last_slice_to_reject` frames of the shower videos will be discarded when 

32 loading the waveforms. (This is required at the moment for LST data) 

33 """ 

34 

35 def __init__( 

36 self, 

37 R0_files: Path | str | Iterable[Path] | Iterable[str], 

38 batch_size: int, 

39 nb_first_slice_to_reject: int = None, 

40 nb_last_slice_to_reject: int = None, 

41 ): 

42 if isinstance(R0_files, (str, Path)): # path to folder containing R0.h5 files or to a single file. 

43 R0_files = Path(R0_files) 

44 if R0_files.is_file(): 

45 R0_files = [R0_files] 

46 elif R0_files.is_dir(): 

47 R0_files = list(R0_files.glob("*.h5")) 

48 else: 

49 raise FileNotFoundError(f"{R0_files} is not a valid file or directory.") 

50 elif isinstance(R0_files, Iterable): # assume R0_files is list of paths to files directly 50 ↛ 53line 50 didn't jump to line 53 because the condition on line 50 was always true

51 R0_files = [Path(p) for p in R0_files] 

52 

53 self._files = sorted(R0_files) # sort list of files for consistency among different execution. 

54 # Compute the number of shower events (and cumulative) in each files 

55 self._files_lengths = [] 

56 for f in self._files: 

57 with tables.open_file(str(f), "r") as f_reader: 

58 self._files_lengths.append(f_reader.root.r0.event.telescope.waveform.tel_001.shape[0]) 

59 self._cum_files_lengths = np.cumsum(self._files_lengths) 

60 self._batch_size = batch_size 

61 self._nb_first_slice_to_reject = nb_first_slice_to_reject 

62 self._nb_last_slice_to_reject = nb_last_slice_to_reject 

63 

64 def read_gains(self) -> np.ndarray: 

65 """Read the per-pixel gains corresponding to the waveforms data. 

66 

67 The high and low gains are stored thus read together. 

68 

69 Returns 

70 ------- 

71 np.ndarray 

72 The per-pixel gains values. gains[0, :] are the high gains, gains[1, :] the low gains. 

73 Shape: (2, N_pixels). 

74 

75 Notes 

76 ----- 

77 At the moment, the gains are constant for all events of a telescope run. Therefore the gains are 

78 simply read from the first file. 

79 """ 

80 with tables.open_file(str(self._files[0]), "r") as f_reader: 

81 gains = f_reader.root.r0.monitoring.telescope.gain.tel_001[:] 

82 return gains 

83 

84 def read_pedestals(self) -> np.ndarray: 

85 """Read the per-pixel pedestals corresponding to the waveform data. 

86 

87 The pedestals corresponding to each channel (gain) are stored and read stacked together. 

88 

89 Returns 

90 ------- 

91 np.ndarray 

92 The per-pixel pedestal values. pedestals[0, :] are the high gain pedestals, while 

93 pedestals[1, :] are the low gain pedestals. Shape: (2, N_pixels) 

94 

95 Notes 

96 ----- 

97 At the moment, the pedestals are constant for all events of a telescope run. Therefore the 

98 pedestals are simply read from the first file. 

99 """ 

100 with tables.open_file(str(self._files[0]), "r") as f_reader: 

101 pedestals = f_reader.root.r0.monitoring.telescope.pedestal.tel_001.cols.pedestal[:].squeeze() 

102 return pedestals 

103 

104 def read_camera_geometry(self) -> np.ndarray: 

105 """Read the camera geometry (pixel coordinates) from the first file of the dataset. 

106 

107 Returns 

108 ------- 

109 np.ndarray 

110 Camera pixel coordinates. geometry[0, :] are the x coordinates, geometry[1, :] are the y 

111 coordinates. Shape (2, N_pixels) 

112 """ 

113 with tables.open_file(str(self._files[0]), "r") as f_reader: 

114 return np.stack( 

115 [ 

116 f_reader.root.configuration.instrument.telescope.camera.geometry_LSTCam.cols.pix_x[:], 

117 f_reader.root.configuration.instrument.telescope.camera.geometry_LSTCam.cols.pix_y[:], 

118 ], 

119 axis=0, 

120 ) 

121 

122 def read_reference_pulse(self) -> tuple[np.ndarray, np.ndarray]: 

123 """Read the reference pulse shape arrays from the first file of the dataset. 

124 

125 Returns 

126 ------- 

127 Tuple[np.ndarray, np.ndarray] 

128 The first element are the "reference_pulse_sample_time", the second are the 

129 "reference_pulse_shape_channel0" and "reference_pulse_shape_channel0", stacked in a single array. 

130 reference_pulse_shape[0] si channel 0 and reference_pulse_shape[1] is channel 1. 

131 """ 

132 with tables.open_file(str(self._files[0]), "r") as f_reader: 

133 ref_pulse = f_reader.root.configuration.instrument.telescope.camera.readout_LSTCam[:] 

134 return ( 

135 ref_pulse["reference_pulse_sample_time"], 

136 np.stack( 

137 [ref_pulse["reference_pulse_shape_channel0"], ref_pulse["reference_pulse_shape_channel1"]], axis=0 

138 ), 

139 ) 

140 

141 def events_n_frames(self) -> int: 

142 """Return the number of frames in an event by reading the 1st event of 1st batch.""" 

143 return self[0][0].shape[1] 

144 

145 def __len__(self) -> int: 

146 """Return the number of batches in the dataset.""" 

147 return np.ceil(self._cum_files_lengths[-1] / self._batch_size).astype(np.int64) 

148 

149 def __getitem__(self, idx) -> tuple[np.ndarray, np.ndarray]: 

150 """Load a batch of waveforms from the dataset. 

151 

152 The size of the batch should be `batch_size` unless there are not enough events left in the dataset 

153 to load (for instance for the last batch, or if the batch size is greater than the number of events 

154 in the dataset.) 

155 

156 The loaded waveforms are truncated at each ends by `nb_first_slice_to_reject`, respectively 

157 `nb_last_slice_to_reject`. 

158 

159 Parameters 

160 ---------- 

161 idx : int 

162 The index of the batch to load. 

163 

164 Returns 

165 ------- 

166 Tuple[np.ndarray, np.ndarray] 

167 waveforms[0] are the high gain waveforms of the batch, waveforms[1] are the low gain waveforms. 

168 Shape of the waveforms arrays: (N_batch, N_frames, N_pixels) 

169 """ 

170 if idx >= len(self) or idx < 0: 

171 raise IndexError(f"Index {idx} out of range for dataset with length {len(self)}") 

172 # Get the index of the file into which the first event should be read. 

173 file_idx = np.searchsorted(self._cum_files_lengths, idx * self._batch_size) 

174 

175 # Calculate the idx of the first event to load in the first file: 

176 # if this is the first file in the dataset, simply subtract the number of already read events 

177 # if this is another file, also subtract the number of events of the previous files 

178 idx_start = ( 

179 idx * self._batch_size - self._cum_files_lengths[file_idx - 1] if file_idx > 0 else idx * self._batch_size 

180 ) 

181 

182 # Load from idx start to idx start + batch size 

183 # if there are not enough events in files, we will get a array smaller than batch size 

184 # and we can continue loading from the next file. 

185 with tables.open_file(self._files[file_idx], "r") as f_reader: 

186 waveforms_high = [ 

187 f_reader.root.r0.event.telescope.waveform.tel_001.cols.waveformHi[ 

188 idx_start : idx_start + self._batch_size 

189 ] 

190 ] 

191 waveforms_low = [ 

192 f_reader.root.r0.event.telescope.waveform.tel_001.cols.waveformLo[ 

193 idx_start : idx_start + self._batch_size 

194 ] 

195 ] 

196 

197 # While we haven't batch_size events, we loop over the next files to load more 

198 # if we are out of files, return the events we have. 

199 cum_lengths = waveforms_high[0].shape[0] 

200 file_idx += 1 

201 while cum_lengths < self._batch_size and file_idx < len(self._files): 

202 with tables.open_file(self._files[file_idx], "r") as f_reader: 

203 waveforms_high.append( 

204 f_reader.root.r0.event.telescope.waveform.tel_001.cols.waveformHi[ 

205 0 : self._batch_size - cum_lengths 

206 ] 

207 ) 

208 waveforms_low.append( 

209 f_reader.root.r0.event.telescope.waveform.tel_001.cols.waveformLo[ 

210 0 : self._batch_size - cum_lengths 

211 ] 

212 ) 

213 cum_lengths = cum_lengths + waveforms_high[-1].shape[0] 

214 file_idx += 1 

215 

216 # concatenate lists of batches of events into single arrays 

217 waveforms_high = np.concatenate(waveforms_high, axis=0, dtype=np.float32) 

218 waveforms_low = np.concatenate(waveforms_low, axis=0, dtype=np.float32) 

219 

220 # remove the unwanted samples (required for LST data) and return the high gain and low gain waveforms 

221 end_frame = None if self._nb_last_slice_to_reject is None else -self._nb_last_slice_to_reject 

222 return ( 

223 waveforms_high[:, self._nb_first_slice_to_reject : end_frame], 

224 waveforms_low[:, self._nb_first_slice_to_reject : end_frame], 

225 )