Coverage for python/lsst/analysis/tools/actions/plot/barPlots.py: 15%

153 statements  

« prev     ^ index     » next       coverage.py v7.2.5, created at 2023-05-04 11:09 +0000

1# This file is part of analysis_tools. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (https://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# This program is free software: you can redistribute it and/or modify 

10# it under the terms of the GNU General Public License as published by 

11# the Free Software Foundation, either version 3 of the License, or 

12# (at your option) any later version. 

13# 

14# This program is distributed in the hope that it will be useful, 

15# but WITHOUT ANY WARRANTY; without even the implied warranty of 

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

17# GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License 

20# along with this program. If not, see <https://www.gnu.org/licenses/>. 

21from __future__ import annotations 

22 

23__all__ = ("BarPanel", "BarPlot") 

24 

25import operator as op 

26from collections import defaultdict 

27from typing import Mapping 

28 

29import matplotlib.pyplot as plt 

30import numpy as np 

31from lsst.analysis.tools.actions.plot.plotUtils import addPlotInfo 

32from lsst.analysis.tools.interfaces import KeyedData, KeyedDataSchema, PlotAction, Vector 

33from lsst.pex.config import Config, ConfigDictField, DictField, Field 

34from matplotlib.figure import Figure 

35from matplotlib.gridspec import GridSpec 

36from matplotlib.patches import Rectangle 

37 

38 

39class BarPanel(Config): 

40 """A configurable class describing a panel in a bar plot.""" 

41 

42 label = Field[str]( 

43 doc="Panel x-axis label.", 

44 default="label", 

45 ) 

46 bars = DictField[str, str]( 

47 doc="A dict specifying the bar graphs to be plotted in this panel. Keys are used to identify " 

48 "bar graph IDs. Values are used to add to the legend label displayed in the upper corner of the " 

49 "panel.", 

50 optional=False, 

51 ) 

52 yscale = Field[str]( 

53 doc="Y axis scaling.", 

54 default="linear", 

55 ) 

56 

57 

58class BarPlot(PlotAction): 

59 """A plotting tool which can take multiple keyed data inputs 

60 and can create one or more bar graphs. 

61 """ 

62 

63 panels = ConfigDictField( 

64 doc="A configurable dict describing the panels to be plotted, and the bar graphs for each panel.", 

65 keytype=str, 

66 itemtype=BarPanel, 

67 default={}, 

68 ) 

69 cmap = Field[str]( 

70 doc="Color map used for bar lines. All types available via `plt.cm` may be used. " 

71 "A number of custom color maps are also defined: `newtab10`, `bright`, `vibrant`.", 

72 default="newtab10", 

73 ) 

74 

75 def getInputSchema(self) -> KeyedDataSchema: 

76 for panel in self.panels: # type: ignore 

77 for barData in self.panels[panel].bars.items(): # type: ignore 

78 yield barData, Vector 

79 

80 def __call__(self, data: KeyedData, **kwargs) -> Mapping[str, Figure] | Figure: 

81 return self.makePlot(data, **kwargs) 

82 

83 def makePlot( 

84 self, data: KeyedData, plotInfo: Mapping[str, str] = None, **kwargs # type: ignore 

85 ) -> Figure: 

86 """Make an N-panel plot with a user-configurable number of bar graphs 

87 displayed in each panel. 

88 

89 Parameters 

90 ---------- 

91 data : `KeyedData` 

92 The catalog to plot the points from. 

93 plotInfo : `dict` 

94 An optional dictionary of information about the data being 

95 plotted with keys: 

96 

97 `"run"` 

98 Output run for the plots (`str`). 

99 `"tractTableType"` 

100 Table from which results are taken (`str`). 

101 `"plotName"` 

102 Output plot name (`str`) 

103 `"SN"` 

104 The global signal-to-noise data threshold (`float`) 

105 `"skymap"` 

106 The type of skymap used for the data (`str`). 

107 `"tract"` 

108 The tract that the data comes from (`int`). 

109 `"bands"` 

110 The bands used for this data (`str` or `list`). 

111 `"visit"` 

112 The visit that the data comes from (`int`) 

113 

114 Returns 

115 ------- 

116 fig : `matplotlib.figure.Figure` 

117 The resulting figure. 

118 

119 """ 

120 

121 # set up figure 

122 fig = plt.figure(dpi=400) 

123 bar_fig, side_fig = fig.subfigures(1, 2, wspace=0, width_ratios=[3, 1]) 

124 axs = self._makeAxes(bar_fig) 

125 

126 # loop over each panel; plot bar graphs 

127 cols = self._assignColors() 

128 all_handles, all_nums, all_vector_labels, all_x_values = [], [], [], [] 

129 for panel, ax in zip(self.panels, axs): 

130 nums, sorted_label, sorted_x_values = self._makePanel(data, panel, ax, cols[panel], **kwargs) 

131 handles, labels = ax.get_legend_handles_labels() # code for plotting 

132 all_handles += handles 

133 all_nums += nums 

134 all_vector_labels += sorted_label 

135 all_x_values += sorted_x_values 

136 

137 # add side panel; add statistics 

138 self._addStatisticsPanel(side_fig, all_handles, all_nums, all_vector_labels, all_x_values) 

139 

140 # add general plot info 

141 if plotInfo is not None: 

142 bar_fig = addPlotInfo(bar_fig, plotInfo) 

143 

144 # finish up 

145 bar_fig.text(0.01, 0.42, "Frequency", rotation=90, transform=bar_fig.transFigure) 

146 plt.draw() 

147 return fig 

148 

149 def _makeAxes(self, fig): 

150 """Determine axes layout for main bar graph figure.""" 

151 num_panels = len(self.panels) 

152 if num_panels <= 1: 

153 ncols = 1 

154 else: 

155 ncols = 2 

156 nrows = int(np.ceil(num_panels / ncols)) 

157 

158 gs = GridSpec(nrows, ncols, left=0.13, right=0.99, bottom=0.1, top=0.88, wspace=0.25, hspace=0.45) 

159 

160 axs = [] 

161 counter = 0 

162 for row in range(nrows): 

163 for col in range(ncols): 

164 counter += 1 

165 if counter < num_panels: 

166 axs.append(fig.add_subplot(gs[row : row + 1, col : col + 1])) 

167 else: 

168 axs.append(fig.add_subplot(gs[row : row + 1, col : np.min([col + 2, ncols + 1])])) 

169 break 

170 

171 return axs 

172 

173 def _assignColors(self): 

174 """Assign colors to bar graphs using a given color map.""" 

175 custom_cmaps = dict( 

176 # https://www.tableau.com/about/blog/2016/7/colors-upgrade-tableau-10-56782 

177 newtab10=[ 

178 "#4e79a7", 

179 "#f28e2b", 

180 "#e15759", 

181 "#76b7b2", 

182 "#59a14f", 

183 "#edc948", 

184 "#b07aa1", 

185 "#ff9da7", 

186 "#9c755f", 

187 "#bab0ac", 

188 ], 

189 # https://personal.sron.nl/~pault/#fig:scheme_bright 

190 bright=[ 

191 "#4477AA", 

192 "#EE6677", 

193 "#228833", 

194 "#CCBB44", 

195 "#66CCEE", 

196 "#AA3377", 

197 "#BBBBBB", 

198 ], 

199 # https://personal.sron.nl/~pault/#fig:scheme_vibrant 

200 vibrant=[ 

201 "#EE7733", 

202 "#0077BB", 

203 "#33BBEE", 

204 "#EE3377", 

205 "#CC3311", 

206 "#009988", 

207 "#BBBBBB", 

208 ], 

209 ) 

210 if self.cmap in custom_cmaps.keys(): 

211 all_cols = custom_cmaps[self.cmap] 

212 else: 

213 try: 

214 all_cols = getattr(plt.cm, self.cmap).copy().colors 

215 except AttributeError: 

216 raise ValueError(f"Unrecognized color map: {self.cmap}") 

217 

218 counter = 0 

219 cols = defaultdict(list) 

220 for panel in self.panels: 

221 for bar in self.panels[panel].bars: 

222 cols[panel].append(all_cols[counter % len(all_cols)]) 

223 counter += 1 

224 return cols 

225 

226 def _makePanel(self, data, panel, ax, col, **kwargs): 

227 """Plot a single panel containing bar graphs.""" 

228 nums = [] 

229 x_values, assigned_labels, assigned_colors = self._assignBinElements(data, panel, col) 

230 sorted_x_values, sorted_labels, sorted_colors = self._sortBarBins( 

231 x_values, assigned_labels, assigned_colors 

232 ) 

233 width, columns = self._getBarWidths(sorted_x_values) 

234 

235 for i, bin in enumerate(sorted_x_values): 

236 bar_data = op.countOf(data[sorted_labels[i]][np.isfinite(data[sorted_labels[i]])], bin) 

237 

238 if width[i] == 1: 

239 bin_center = bin 

240 else: 

241 bin_center = bin - 0.35 + width[i] * columns[i] 

242 

243 ax.bar(bin_center, bar_data, width[i], lw=2, label=sorted_labels[i], color=sorted_colors[i]) 

244 nums.append(bar_data) 

245 

246 # Get plot range 

247 x_range = [x for x in range(int(min(sorted_x_values)), int(max(sorted_x_values)) + 1)] 

248 ax.set_xticks(x_range) 

249 ax.set_xlabel(self.panels[panel].label) 

250 ax.set_yscale(self.panels[panel].yscale) 

251 ax.tick_params(labelsize=7) 

252 # add a buffer to the top of the plot to allow headspace for labels 

253 ylims = list(ax.get_ylim()) 

254 if ax.get_yscale() == "log": 

255 ylims[1] = 10 ** (np.log10(ylims[1]) * 1.1) 

256 else: 

257 ylims[1] *= 1.1 

258 ax.set_ylim(ylims[0], ylims[1]) 

259 return nums, sorted_labels, sorted_x_values 

260 

261 def _assignBinElements(self, data, panel, col): 

262 labels = [] 

263 assigned_labels = [] 

264 x_values = [] 

265 assigned_colors = [] 

266 n_labels = 0 

267 

268 for bar in self.panels[panel].bars: 

269 labels.append(bar) 

270 n_labels += 1 

271 

272 # If a label has multiple unique elements in it, repeats the label 

273 i = 0 

274 for single_label in labels: 

275 unique_elements = np.unique(data[single_label]) 

276 

277 for bin in unique_elements: 

278 x_values.append(int(bin)) 

279 

280 for count in range(len(unique_elements)): 

281 assigned_labels.append(single_label) 

282 assigned_colors.append(col[i]) # Assign color from color cmap 

283 

284 i += 1 

285 

286 return x_values, assigned_labels, assigned_colors 

287 

288 def _sortBarBins(self, x_values, assigned_labels, assigned_colors): 

289 """Sorts the existing x_values, assigned_labels, 

290 and assigned_colors/x_value from lowest to 

291 highest and then uses the sorted indices to sort 

292 all x, labels, and colors in that order. 

293 """ 

294 

295 sorted_indices = np.argsort(x_values) 

296 

297 sorted_labels = [] 

298 sorted_x_values = [] 

299 sorted_colors = [] 

300 

301 for position in sorted_indices: 

302 sorted_x_values.append(x_values[position]) 

303 sorted_labels.append(assigned_labels[position]) 

304 sorted_colors.append(assigned_colors[position]) 

305 

306 return sorted_x_values, sorted_labels, sorted_colors 

307 

308 def _getBarWidths(self, x_values): 

309 """Determine the width of the panels in each 

310 bin and which column is assigned.""" 

311 width = [] 

312 columns = [] 

313 current_column = 0 

314 current_i = 0 

315 

316 for i in x_values: 

317 # Number of repeating values 

318 n_repeating = x_values.count(i) 

319 width.append(1.0 / n_repeating) 

320 if n_repeating > 1 and current_column != 0 and current_i == i: 

321 columns.append(current_column) 

322 current_column += 1 

323 

324 else: 

325 current_column = 0 

326 columns.append(current_column) 

327 current_i = i 

328 current_column += 1 

329 

330 return width, columns 

331 

332 def _addStatisticsPanel(self, fig, handles, nums, sorted_labels, sorted_x_value): 

333 """Add an adjoining panel containing bar graph summary statistics.""" 

334 ax = fig.add_subplot(1, 1, 1) 

335 ax.axis("off") 

336 plt.subplots_adjust(left=0.05, right=0.95, bottom=0.09, top=0.9) 

337 

338 # empty handle, used to populate the bespoke legend layout 

339 empty = Rectangle((0, 0), 1, 1, fc="w", fill=False, edgecolor="none", linewidth=0) 

340 

341 # set up new legend handles and labels 

342 

343 legend_handles = [empty] + handles + ([empty] * 3 * len(handles)) + ([empty] * 3) 

344 legend_labels = ( 

345 ([""] * (len(handles) + 1)) 

346 + ["Bin"] 

347 + sorted_x_value 

348 + ["Count"] 

349 + nums 

350 + ["Sources"] 

351 + sorted_labels 

352 ) 

353 

354 # add the legend 

355 ax.legend( 

356 legend_handles, 

357 legend_labels, 

358 loc="lower left", 

359 ncol=4, 

360 handletextpad=-0.25, 

361 fontsize=6, 

362 borderpad=0, 

363 frameon=False, 

364 columnspacing=-0.25, 

365 )