Coverage for pyhiperta/calibration.py: 100%

1"""Implements the "calibration" of R0 data: going from 2 gain ADC counts to number of photo-electrons""" 

2 

3from typing import Tuple 

4 

5import numpy as np 

6 

7 

8def select_channel( 

9 waveform_high: np.ndarray, 

10 waveform_low: np.ndarray, 

11 gains: np.ndarray, 

12 pedestals: np.ndarray, 

13 window_integration_correction: np.ndarray, 

14 threshold: float, 

15) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: 

16 """Select values to use in waveforms high/low gain and corresponding gains and pedestals based on threshold. 

17 

18 The choice is independent per shower and per pixel, but shared in time for a given pixel: 

19 For each pixel, if any value of `waveform_high` is above `threshold` then the low gain and 

20 corresponding waveform channel and pedestal is chosen. Otherwise the high gain is chosen. 

21 

22 Parameters 

23 ---------- 

24 waveforms_high : np.ndarray 

25 R0 waveform high gain. Shape: (N_batch, N_frames, N_pixels) 

26 waveform_low : np.ndarray 

27 R0 waveform low gain. Shape: (N_batch, N_frames, N_pixels) 

28 gains : np.ndarray 

29 Per-pixel high and low gains. `gains[0]` is high gain, `gains[1]` is low gains. shape (2, N_pixels) 

30 pedestals : np.ndarray 

31 Per-pixel pedestal for low and high gains. `pedestals[0]` corresponds to high gain, `pedestals[1]` 

32 corresponds to low gain. Shape (2, N_pixels) 

33 window_integration_correction : np.ndarray 

34 Array of shape (2,) containing the correction to apply after a windowed integration, for each gain. 

35 threshold : float 

36 Threshold to chose the waveform channel. A value above `threshold` indicates that the high gain waveform 

37 saturated and low gain should be used. 

38 

39 Returns 

40 ------- 

41 Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray] 

42 The waveform selected from each gains based on `threshold`, associated gains, pedestals and integration 

43 correction to used. The order is: 

44 - waveform : np.ndarray 

45 waveform where the appropriate gain has been selected. Shape (N_batch, N_frames, N_pixels) 

46 - gains : np.ndarray 

47 Gain to use to calibrate `waveform` (high gain when waveform_high wasn't above `threshold` and 

48 low gain otherwise). Shape (N_batch, N_pixels) 

49 - pedestals : np.ndarray 

50 Pedestals to use to calibrate `waveform`. Shape (N_batch, N_pixels) 

51 - window_integration_correction : np.ndarray 

52 Correction to use after the integration. Shape(N_batch, N_pixels) 

53 """ 

54 idx_saturated = (waveform_high > threshold).any(axis=-2, keepdims=True) 

55 return ( 

56 np.where(idx_saturated, waveform_low, waveform_high), 

57 np.where(np.squeeze(idx_saturated, -2), gains[1], gains[0]), 

58 np.where(np.squeeze(idx_saturated, -2), pedestals[1], pedestals[0]), 

59 np.where(np.squeeze(idx_saturated, -2), window_integration_correction[1], window_integration_correction[0]), 

60 ) 

61 

62 

63def calibrate(waveform: np.ndarray, gains: np.ndarray, pedestals: np.ndarray) -> np.ndarray: 

64 """Apply pedestal and gain calibration to waveforms. 

65 

66 Parameters 

67 ---------- 

68 waveform : np.ndarray 

69 Batch or single video samples of the shower events. Shape: (N_batch, N_frames, N_pixels) 

70 gains : np.ndarray 

71 For each waveform, the gains to apply to each pixel. The same gain is applied for all 

72 frames of a shower pixel. Shape: (N_batch, N_pixels). 

73 pedestals : np.ndarray 

74 For each waveform, the pedestals to apply to each pixel. The same pedestal is applied for 

75 all frames of a shower pixel. Shape (N_batch, N_pixels) 

76 

77 Returns 

78 ------- 

79 np.ndarray 

80 Calibrated waveforms (number of photo-electrons). Shape is identical to `waveform`. 

81 """ 

82 return (waveform - pedestals[..., np.newaxis, :]) * gains[..., np.newaxis, :]