Coverage for tcvx21/plotting/labels_m.py: 87%

1""" 

2Routines for adding labels to plots 

3""" 

4import matplotlib.pyplot as plt 

5import numpy as np 

6from matplotlib.offsetbox import AnchoredText 

7 

8from tcvx21.units_m import Quantity 

9from tcvx21.observable_c.observable_m import Observable 

10 

11# Locations for legend or anchored text 

12locations = {'best': 0, 'upper right': 1, 'upper left': 2, 'lower left': 3, 

13 'lower right': 4, 'right': 5, 'center left': 6, 'center right': 7, 

14 'lower center': 8, 'upper center': 9, 'center': 10} 

15 

16 

17def add_y_zero_line(ax, level=0.0): 

18 """Mark the 0-line""" 

19 ax.axhline(level, color='k', linestyle='-', linewidth=0.5) 

20 

21 

22def add_x_zero_line(ax): 

23 """Mark the separatrix""" 

24 ax.axvline(0.0, color='k', linestyle='-', linewidth=0.5) 

25 

26 

27def add_twinx_label(ax, text: str, labelpad=20, visible=True): 

28 """Adds a right-hand-side label with the field direction""" 

29 new_ax = ax.twinx() 

30 new_ax.set_yticks([]) 

31 new_ax.set_ylabel(text, rotation=270, labelpad=labelpad) 

32 

33 if not visible: 

34 new_ax.set_xticks([]) 

35 new_ax.spines['top'].set_visible(False) 

36 new_ax.spines['right'].set_visible(False) 

37 new_ax.spines['bottom'].set_visible(False) 

38 new_ax.spines['left'].set_visible(False) 

39 

40 # Switch focus back to the main axis 

41 plt.sca(ax) 

42 

43 

44def format_yaxis(ax): 

45 """ 

46 Move the exponent into the ylabel 

47 You have to call plt.draw() before this, otherwise the offset will be blank 

48 """ 

49 plt.draw() 

50 ax.yaxis.offsetText.set_visible(False) 

51 units_string = f"{Quantity(1, ax.yaxis.units).units:~P}" if ax.yaxis.units else "Unitless" 

52 

53 units_string = units_string.lstrip('1') 

54 

55 # Get the axis scale factor (i.e. 10^19 for the density) 

56 scale_factor = ax.get_yaxis().get_offset_text().get_text() 

57 

58 label_string = f"$10^{{{scale_factor.split('e')[-1]}}}${units_string}" if scale_factor \ 

59 else f"{units_string}" 

60 

61 ax.set_ylabel(label_string) 

62 

63 

64def make_colorbar(ax, mappable, units, as_title: bool = True, 

65 remove_offset: bool = True): 

66 """ 

67 Adds a colorbar to the image, with a units label 

68 """ 

69 cbar = plt.colorbar(mappable, ax=ax) 

70 cax = cbar.ax 

71 

72 plt.draw() 

73 units_string = f"{Quantity(1, units).units:~P}" if units else "" 

74 

75 if remove_offset: 

76 # If the units string starts with a '1' (i.e. 1/m^3), remove the 1 

77 units_string = units_string.lstrip('1') 

78 cax.yaxis.offsetText.set_visible(False) 

79 

80 # Get the axis scale factor (i.e. 10^19 for the density) 

81 scale_factor = cax.get_yaxis().get_offset_text().get_text() 

82 

83 label_string = f"[$10^{{{scale_factor.split('e')[-1]}}}${units_string}]" if scale_factor \ 

84 else f"[{units_string}]" 

85 else: 

86 label_string = f"[{units_string}]" 

87 

88 if as_title: 

89 cax.set_title(label_string, fontsize=plt.rcParams['ytick.labelsize']) 

90 else: 

91 cax.set_ylabel(label_string, rotation=270, labelpad=15) 

92 

93 return cbar 

94 

95 

96def make_field_direction_string(field_direction: str): 

97 """Returns the field direction as a nice string""" 

98 field_direction_rename = { 

99 'forward_field': 'Forward field', 

100 'reversed_field': 'Reversed field', 

101 } 

102 return field_direction_rename[field_direction] 

103 

104 

105def make_diagnostic_string(diagnostic: str): 

106 """Returns the diagnostic as a nice string""" 

107 diagnostic_rename = { 

108 'LFS-LP': 'Low-field-side target', 

109 'LFS-IR': 'Low-field-side target', 

110 'HFS-LP': 'High-field-side target', 

111 'TS': 'Divertor Thomson', 

112 'RDPA': 'Divertor volume', 

113 'FHRP': 'Outboard midplane' 

114 } 

115 return diagnostic_rename[diagnostic] 

116 

117 

118def make_observable_string(observable: Observable): 

119 """Returns the observable as a compact string""" 

120 observable_string = observable.name 

121 

122 shorten_title_rename = { 

123 'Standard deviation': 'Std. dev.', 

124 'of the': 'of', 

125 'Pearson kurtosis': 'Kurtosis', 

126 'Plasma velocity normalised to local sound speed': 'Mach number', 

127 'Ion saturation current': '$J_{sat}$', 

128 'ion saturation current': '$J_{sat}$' 

129 } 

130 for key, value in shorten_title_rename.items(): 

131 observable_string = observable_string.replace(key, value) 

132 

133 return observable_string 

134 

135 

136def make_labels(field_direction: str, diagnostic: str, observable: Observable): 

137 """ 

138 Returns strings which can be used in a title 

139 """ 

140 

141 return make_field_direction_string(field_direction), \ 

142 make_diagnostic_string(diagnostic), \ 

143 make_observable_string(observable) 

144 

145 

146def label_subplots(axs: np.ndarray, prefix: str = None): 

147 """Add subplot labels to an array of subplots""" 

148 plt.draw() 

149 

150 index = -1 

151 

152 for ax in axs.flatten(): 

153 

154 if not ax.get_visible(): 

155 continue 

156 index += 1 

157 

158 # Convert integer to character 

159 if prefix is None: 

160 label = f"{chr(index + 65)}" 

161 else: 

162 label = f"{prefix}{chr(index + 65)}" 

163 

164 annotation_kwargs = dict(xycoords="axes fraction", 

165 bbox=dict( 

166 boxstyle="round", 

167 alpha=0.9, 

168 facecolor="white", 

169 edgecolor="grey" 

170 )) 

171 

172 annotate = ax.annotate(label, xy=(0.05, 0.85), **annotation_kwargs) 

173 

174 # Try to not overlap with the legend 

175 legend = ax.get_legend() 

176 try: 

177 legend_extent = legend.get_window_extent() 

178 except AttributeError: 

179 continue 

180 annotate_extent = annotate.get_window_extent() 

181 

182 if (((legend_extent.x0 <= annotate_extent.x0 <= legend_extent.x1) or 

183 (legend_extent.x0 <= annotate_extent.x1 <= legend_extent.x1)) and 

184 ((legend_extent.y0 <= annotate_extent.y0 <= legend_extent.y1) or 

185 (legend_extent.y0 <= annotate_extent.y1 <= legend_extent.y1))): 

186 # Prevent annotation from overlapping with legend 

187 

188 annotate.remove() 

189 ax.annotate(label, xy=(0.05, 0.1), **annotation_kwargs)