Coverage for polars_analysis / plotting / cross_talk_plotting.py: 92%

1import logging 

2import pathlib 

3from itertools import product 

4 

5import matplotlib.pyplot as plt 

6import numpy as np 

7import polars as pl 

8 

9from polars_analysis.plotting import helper 

10 

11# Instantiate logger 

12log = logging.getLogger(__name__) 

13 

14""" 

15Functions for plotting pedestal cross-talk data, which has already been processed. 

16""" 

17 

18 

19# --- 

20def _find_gain_ratio(df: pl.DataFrame) -> float: 

21 return ( 

22 df.filter(pl.col("is_reference_pulse")) 

23 .select((pl.col("mean_interleaved_pulse").list.max()).sort()) 

24 .select(pl.all() / pl.all().first()) 

25 .tail(1) 

26 .item() 

27 ) 

28 

29 

30# --- 

31def plot_ct_overlay(df: pl.DataFrame, plot_dir: pathlib.Path) -> None: 

32 atten_val: int = int(df[0]["att_val"][0]) 

33 board_id: str = df[0]["board_id"][0] 

34 run_num: int = int(df[0]["run_number"][0]) 

35 

36 gr = _find_gain_ratio(df) 

37 

38 # Gets reference (max amplitude) pulse 

39 samples: str = "mean_interleaved_pulse" 

40 reference_samples = df.sort(pl.col(samples).list.max()).select(pl.last(samples)).item().to_numpy() 

41 

42 left_bound = np.argmax(reference_samples) - 200 

43 right_bound = np.argmax(reference_samples) + 1000 

44 pulse_channel = df.filter(pl.col("is_reference_pulse")).select("channel").unique().item() 

45 

46 _, ax = plt.subplots() 

47 

48 for channel in df["channel"].unique(): 

49 try: 

50 data_lo = ( 

51 df.filter(pl.col("channel") == channel).filter(pl.col("gain") == "lo").select(samples).item().to_numpy() 

52 ) 

53 except ValueError: 

54 log.warning(f"lo gain ch{channel} missing and skipped in overlay") 

55 data_lo = np.array([]) 

56 

57 try: 

58 data_hi = ( 

59 df.filter(pl.col("channel") == channel).filter(pl.col("gain") == "hi").select(samples).item().to_numpy() 

60 ) 

61 except ValueError: 

62 log.warning(f"hi gain ch{channel} missing and skipped in overlay") 

63 data_hi = [] 

64 

65 label = f"Ch {channel}" 

66 

67 if channel == pulse_channel: 

68 if len(data_lo) > 0: 

69 ax.plot(gr * data_lo[left_bound:right_bound], label=label + "L * GR", zorder=2) 

70 if len(data_hi) > 0: 

71 ax.plot(data_hi[left_bound:right_bound], label=label + "H", zorder=2) 

72 else: 

73 if len(data_lo) > 0: 

74 ax.plot(100 * gr * data_lo[left_bound:right_bound], label=label + "L * GR * 100", zorder=0) 

75 if len(data_hi) > 0: 

76 ax.plot(100 * data_hi[left_bound:right_bound], label=label + "H * 100", zorder=1) 

77 

78 ax.plot([], [], " ", label=f"GR={gr:.2f}") 

79 ax.set( 

80 xlabel="Time [ns]", 

81 ylabel="Scaled Amplitude [ADC Counts]", 

82 ) 

83 ax.grid() 

84 ax.legend() 

85 plt.title( 

86 helper.plot_summary_string( 

87 name="Crosstalk", board_id=board_id, run_numbers=run_num, channels=pulse_channel, attenuation=atten_val 

88 ) 

89 ) 

90 plt.tight_layout() 

91 

92 save_name = "crosstalk_overlay" 

93 plt.savefig(plot_dir / f"{save_name}.png") 

94 plt.cla() 

95 plt.clf() 

96 plt.close() 

97 return 

98 

99 

100# --- 

101def plot_ct_table(df: pl.DataFrame, plot_dir: pathlib.Path) -> None: 

102 channels = sorted(df["channel"].unique().to_list()) 

103 pulse_channel = df.filter(pl.col("is_reference_pulse")).select("channel").unique().item() 

104 atten_val: int = int(df[0]["att_val"][0]) 

105 board_id: str = df[0]["board_id"][0] 

106 run_num: int = int(df[0]["run_number"][0]) 

107 

108 columns = [] 

109 for ch in channels: 

110 if ch == pulse_channel: 

111 continue 

112 for gain in ["lo", "hi"]: 

113 columns.append(f"channel{ch} {gain}") 

114 

115 rows = [f"{pulse_channel} {gain}" for gain in ["hi"]] 

116 

117 table_data = np.full((len(columns), len(rows)), -99.0) 

118 for x, y in product(range(len(columns)), range(len(rows))): 

119 try: 

120 x_gain = columns[x][-2:] 

121 y_gain = rows[y][-2:] 

122 

123 x_channel = int(columns[x].split()[0].replace("channel", "")) 

124 y_channel = int(rows[y].split()[0]) 

125 

126 x_energy = ( 

127 df.filter((pl.col("channel") == x_channel) & (pl.col("gain") == x_gain)).select("energy_mean").item() 

128 ) 

129 y_energy = ( 

130 df.filter((pl.col("channel") == y_channel) & (pl.col("gain") == y_gain)).select("energy_mean").item() 

131 ) 

132 if y_energy > 0: 

133 table_data[x, y] = 100 * x_energy / y_energy 

134 except ValueError: 

135 log.warning(f"{columns[x]=}, {rows[y]=} missing, skipped in table") 

136 

137 fig, ax = plt.subplots(figsize=(6, 1)) 

138 ax.get_xaxis().set_visible(False) 

139 ax.get_yaxis().set_visible(False) 

140 plt.box(on=None) 

141 ax.axis("tight") 

142 ax.set_title( 

143 helper.plot_summary_string( 

144 name="Crosstalk %", board_id=board_id, run_numbers=run_num, channels=pulse_channel, attenuation=atten_val 

145 ) 

146 ) 

147 str_list = [f"{x:.3f}" for x in list(table_data.transpose()[0])] 

148 arr = np.array(([str_list]), dtype=str) 

149 ax.table(cellText=arr, rowLabels=rows, colLabels=columns, loc="center", bbox=[0, 0, 1, 1]) # type: ignore 

150 fig.tight_layout() 

151 

152 save_name = "crosstalk_table" 

153 plt.savefig(plot_dir / f"{save_name}.png", dpi=300) 

154 log.debug(f"Saved to {plot_dir}/{save_name}.png") 

155 plt.cla() 

156 plt.clf() 

157 plt.close() 

158 return