Coverage for python/lsst/analysis/drp/plotUtils.py: 8%

180 statements  

« prev     ^ index     » next       coverage.py v6.4, created at 2022-05-25 12:45 +0000

1# This file is part of analysis_drp. 

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/>. 

21 

22import numpy as np 

23import matplotlib.pyplot as plt 

24import scipy.odr as scipyODR 

25import matplotlib 

26from matplotlib import colors 

27from typing import List, Tuple 

28 

29from lsst.geom import Box2D, SpherePoint, degrees 

30 

31null_formatter = matplotlib.ticker.NullFormatter() 

32 

33 

34def parsePlotInfo(dataId, runName, tableName, bands, plotName, SN): 

35 """Parse plot info from the dataId 

36 

37 Parameters 

38 ---------- 

39 dataId : `lsst.daf.butler.core.dimensions.` 

40 `_coordinate._ExpandedTupleDataCoordinate` 

41 runName : `str` 

42 

43 Returns 

44 ------- 

45 plotInfo : `dict` 

46 """ 

47 plotInfo = {"run": runName, "tractTableType": tableName, "plotName": plotName, "SN": SN} 

48 

49 for dataInfo in dataId: 

50 plotInfo[dataInfo.name] = dataId[dataInfo.name] 

51 

52 bandStr = "" 

53 for band in bands: 

54 bandStr += (", " + band) 

55 plotInfo["bands"] = bandStr[2:] 

56 

57 if "tract" not in plotInfo.keys(): 

58 plotInfo["tract"] = "N/A" 

59 if "visit" not in plotInfo.keys(): 

60 plotInfo["visit"] = "N/A" 

61 

62 return plotInfo 

63 

64 

65def generateSummaryStats(cat, colName, skymap, plotInfo): 

66 """Generate a summary statistic in each patch or detector 

67 

68 Parameters 

69 ---------- 

70 cat : `pandas.core.frame.DataFrame` 

71 colName : `str` 

72 skymap : `lsst.skymap.ringsSkyMap.RingsSkyMap` 

73 plotInfo : `dict` 

74 

75 Returns 

76 ------- 

77 patchInfoDict : `dict` 

78 """ 

79 

80 # TODO: what is the more generic type of skymap? 

81 tractInfo = skymap.generateTract(plotInfo["tract"]) 

82 tractWcs = tractInfo.getWcs() 

83 

84 if "sourceType" in cat.columns: 

85 cat = cat.loc[cat["sourceType"] != 0] 

86 

87 # For now also convert the gen 2 patchIds to gen 3 

88 

89 patchInfoDict = {} 

90 maxPatchNum = tractInfo.num_patches.x*tractInfo.num_patches.y 

91 patches = np.arange(0, maxPatchNum, 1) 

92 for patch in patches: 

93 if patch is None: 

94 continue 

95 # Once the objectTable_tract catalogues are using gen 3 patches 

96 # this will go away 

97 onPatch = (cat["patch"] == patch) 

98 stat = np.nanmedian(cat[colName].values[onPatch]) 

99 try: 

100 patchTuple = (int(patch.split(",")[0]), int(patch.split(",")[-1])) 

101 patchInfo = tractInfo.getPatchInfo(patchTuple) 

102 gen3PatchId = tractInfo.getSequentialPatchIndex(patchInfo) 

103 except AttributeError: 

104 # For native gen 3 tables the patches don't need converting 

105 # When we are no longer looking at the gen 2 -> gen 3 

106 # converted repos we can tidy this up 

107 gen3PatchId = patch 

108 patchInfo = tractInfo.getPatchInfo(patch) 

109 

110 corners = Box2D(patchInfo.getInnerBBox()).getCorners() 

111 skyCoords = tractWcs.pixelToSky(corners) 

112 

113 patchInfoDict[gen3PatchId] = (skyCoords, stat) 

114 

115 tractCorners = Box2D(tractInfo.getBBox()).getCorners() 

116 skyCoords = tractWcs.pixelToSky(tractCorners) 

117 patchInfoDict["tract"] = (skyCoords, np.nan) 

118 

119 return patchInfoDict 

120 

121 

122def generateSummaryStatsVisit(cat, colName, visitSummaryTable, plotInfo): 

123 """Generate a summary statistic in each patch or detector 

124 

125 Parameters 

126 ---------- 

127 cat : `pandas.core.frame.DataFrame` 

128 colName : `str` 

129 visitSummaryTable : `pandas.core.frame.DataFrame` 

130 plotInfo : `dict` 

131 

132 Returns 

133 ------- 

134 visitInfoDict : `dict` 

135 """ 

136 

137 visitInfoDict = {} 

138 for ccd in cat.detector.unique(): 

139 if ccd is None: 

140 continue 

141 onCcd = (cat["detector"] == ccd) 

142 stat = np.nanmedian(cat[colName].values[onCcd]) 

143 

144 sumRow = (visitSummaryTable["id"] == ccd) 

145 corners = zip(visitSummaryTable["raCorners"][sumRow][0], visitSummaryTable["decCorners"][sumRow][0]) 

146 cornersOut = [] 

147 for (ra, dec) in corners: 

148 corner = SpherePoint(ra, dec, units=degrees) 

149 cornersOut.append(corner) 

150 

151 visitInfoDict[ccd] = (cornersOut, stat) 

152 

153 return visitInfoDict 

154 

155 

156# Inspired by matplotlib.testing.remove_ticks_and_titles 

157def get_and_remove_axis_text(ax) -> Tuple[List[str], List[np.ndarray]]: 

158 """Remove text from an Axis and its children and return with line points. 

159 

160 Parameters 

161 ---------- 

162 ax : `plt.Axis` 

163 A matplotlib figure axis. 

164 

165 Returns 

166 ------- 

167 texts : `List[str]` 

168 A list of all text strings (title and axis/legend/tick labels). 

169 line_xys : `List[numpy.ndarray]` 

170 A list of all line ``_xy`` attributes (arrays of shape ``(N, 2)``). 

171 """ 

172 line_xys = [line._xy for line in ax.lines] 

173 texts = [text.get_text() for text in (ax.title, ax.xaxis.label, ax.yaxis.label)] 

174 ax.set_title("") 

175 ax.set_xlabel("") 

176 ax.set_ylabel("") 

177 

178 try: 

179 texts_legend = ax.get_legend().texts 

180 texts.extend(text.get_text() for text in texts_legend) 

181 for text in texts_legend: 

182 text.set_alpha(0) 

183 except AttributeError: 

184 pass 

185 

186 for idx in range(len(ax.texts)): 

187 texts.append(ax.texts[idx].get_text()) 

188 ax.texts[idx].set_text('') 

189 

190 ax.xaxis.set_major_formatter(null_formatter) 

191 ax.xaxis.set_minor_formatter(null_formatter) 

192 ax.yaxis.set_major_formatter(null_formatter) 

193 ax.yaxis.set_minor_formatter(null_formatter) 

194 try: 

195 ax.zaxis.set_major_formatter(null_formatter) 

196 ax.zaxis.set_minor_formatter(null_formatter) 

197 except AttributeError: 

198 pass 

199 for child in ax.child_axes: 

200 texts_child, lines_child = get_and_remove_axis_text(child) 

201 texts.extend(texts_child) 

202 

203 return texts, line_xys 

204 

205 

206def get_and_remove_figure_text(figure: plt.Figure): 

207 """Remove text from a Figure and its Axes and return with line points. 

208 

209 Parameters 

210 ---------- 

211 figure : `matplotlib.pyplot.Figure` 

212 A matplotlib figure. 

213 

214 Returns 

215 ------- 

216 texts : `List[str]` 

217 A list of all text strings (title and axis/legend/tick labels). 

218 line_xys : `List[numpy.ndarray]`, (N, 2) 

219 A list of all line ``_xy`` attributes (arrays of shape ``(N, 2)``). 

220 """ 

221 texts = [str(figure._suptitle)] 

222 lines = [] 

223 figure.suptitle("") 

224 

225 texts.extend(text.get_text() for text in figure.texts) 

226 figure.texts = [] 

227 

228 for ax in figure.get_axes(): 

229 texts_ax, lines_ax = get_and_remove_axis_text(ax) 

230 texts.extend(texts_ax) 

231 lines.extend(lines_ax) 

232 

233 return texts, lines 

234 

235 

236def addPlotInfo(fig, plotInfo): 

237 """Add useful information to the plot 

238 

239 Parameters 

240 ---------- 

241 fig : `matplotlib.figure.Figure` 

242 plotInfo : `dict` 

243 

244 Returns 

245 ------- 

246 fig : `matplotlib.figure.Figure` 

247 """ 

248 

249 # TO DO: figure out how to get this information 

250 photocalibDataset = "None" 

251 astroDataset = "None" 

252 

253 fig.text(0.01, 0.99, plotInfo["plotName"], fontsize=8, transform=fig.transFigure, ha="left", va="top") 

254 

255 run = plotInfo["run"] 

256 datasetsUsed = f"\nPhotoCalib: {photocalibDataset}, Astrometry: {astroDataset}" 

257 tableType = f"\nTable: {plotInfo['tractTableType']}" 

258 

259 dataIdText = "" 

260 if str(plotInfo["tract"]) != "N/A": 

261 dataIdText += f", Tract: {plotInfo['tract']}" 

262 if str(plotInfo["visit"]) != "N/A": 

263 dataIdText += f", Visit: {plotInfo['visit']}" 

264 

265 bandsText = f", Bands: {''.join(plotInfo['bands'].split(' '))}" 

266 SNText = f", S/N: {plotInfo['SN']}" 

267 infoText = f"\n{run}{datasetsUsed}{tableType}{dataIdText}{bandsText}{SNText}" 

268 fig.text(0.01, 0.98, infoText, fontsize=7, transform=fig.transFigure, alpha=0.6, ha="left", va="top") 

269 

270 return fig 

271 

272 

273def stellarLocusFit(xs, ys, paramDict): 

274 """Make a fit to the stellar locus 

275 

276 Parameters 

277 ---------- 

278 xs : `numpy.ndarray` 

279 The color on the xaxis 

280 ys : `numpy.ndarray` 

281 The color on the yaxis 

282 paramDict : lsst.pex.config.dictField.Dict 

283 A dictionary of parameters for line fitting 

284 xMin : `float` 

285 The minimum x edge of the box to use for initial fitting 

286 xMax : `float` 

287 The maximum x edge of the box to use for initial fitting 

288 yMin : `float` 

289 The minimum y edge of the box to use for initial fitting 

290 yMax : `float` 

291 The maximum y edge of the box to use for initial fitting 

292 mHW : `float` 

293 The hardwired gradient for the fit 

294 bHW : `float` 

295 The hardwired intercept of the fit 

296 

297 Returns 

298 ------- 

299 paramsOut : `dict` 

300 A dictionary of the calculated fit parameters 

301 xMin : `float` 

302 The minimum x edge of the box to use for initial fitting 

303 xMax : `float` 

304 The maximum x edge of the box to use for initial fitting 

305 yMin : `float` 

306 The minimum y edge of the box to use for initial fitting 

307 yMax : `float` 

308 The maximum y edge of the box to use for initial fitting 

309 mHW : `float` 

310 The hardwired gradient for the fit 

311 bHW : `float` 

312 The hardwired intercept of the fit 

313 mODR : `float` 

314 The gradient calculated by the ODR fit 

315 bODR : `float` 

316 The intercept calculated by the ODR fit 

317 yBoxMin : `float` 

318 The y value of the fitted line at xMin 

319 yBoxMax : `float` 

320 The y value of the fitted line at xMax 

321 bPerpMin : `float` 

322 The intercept of the perpendicular line that goes through xMin 

323 bPerpMax : `float` 

324 The intercept of the perpendicular line that goes through xMax 

325 mODR2 : `float` 

326 The gradient from the second round of fitting 

327 bODR2 : `float` 

328 The intercept from the second round of fitting 

329 mPerp : `float` 

330 The gradient of the line perpendicular to the line from the 

331 second fit 

332 

333 Notes 

334 ----- 

335 The code does two rounds of fitting, the first is initiated using the 

336 hardwired values given in the `paramDict` parameter and is done using 

337 an Orthogonal Distance Regression fit to the points defined by the 

338 box of xMin, xMax, yMin and yMax. Once this fitting has been done a 

339 perpendicular bisector is calculated at either end of the line and 

340 only points that fall within these lines are used to recalculate the fit. 

341 """ 

342 

343 # Points to use for the fit 

344 fitPoints = np.where((xs > paramDict["xMin"]) & (xs < paramDict["xMax"]) 

345 & (ys > paramDict["yMin"]) & (ys < paramDict["yMax"]))[0] 

346 

347 linear = scipyODR.polynomial(1) 

348 

349 data = scipyODR.Data(xs[fitPoints], ys[fitPoints]) 

350 odr = scipyODR.ODR(data, linear, beta0=[paramDict["bHW"], paramDict["mHW"]]) 

351 params = odr.run() 

352 mODR = float(params.beta[1]) 

353 bODR = float(params.beta[0]) 

354 

355 paramsOut = {"xMin": paramDict["xMin"], "xMax": paramDict["xMax"], "yMin": paramDict["yMin"], 

356 "yMax": paramDict["yMax"], "mHW": paramDict["mHW"], "bHW": paramDict["bHW"], 

357 "mODR": mODR, "bODR": bODR} 

358 

359 # Having found the initial fit calculate perpendicular ends 

360 mPerp = -1.0/mODR 

361 # When the gradient is really steep we need to use 

362 # the y limits of the box rather than the x ones 

363 

364 if np.abs(mODR) > 1: 

365 yBoxMin = paramDict["yMin"] 

366 xBoxMin = (yBoxMin - bODR)/mODR 

367 yBoxMax = paramDict["yMax"] 

368 xBoxMax = (yBoxMax - bODR)/mODR 

369 else: 

370 yBoxMin = mODR*paramDict["xMin"] + bODR 

371 xBoxMin = paramDict["xMin"] 

372 yBoxMax = mODR*paramDict["xMax"] + bODR 

373 xBoxMax = paramDict["xMax"] 

374 

375 bPerpMin = yBoxMin - mPerp*xBoxMin 

376 

377 paramsOut["yBoxMin"] = yBoxMin 

378 paramsOut["bPerpMin"] = bPerpMin 

379 

380 bPerpMax = yBoxMax - mPerp*xBoxMax 

381 

382 paramsOut["yBoxMax"] = yBoxMax 

383 paramsOut["bPerpMax"] = bPerpMax 

384 

385 # Use these perpendicular lines to chose the data and refit 

386 fitPoints = ((ys > mPerp*xs + bPerpMin) & (ys < mPerp*xs + bPerpMax)) 

387 data = scipyODR.Data(xs[fitPoints], ys[fitPoints]) 

388 odr = scipyODR.ODR(data, linear, beta0=[bODR, mODR]) 

389 params = odr.run() 

390 mODR = float(params.beta[1]) 

391 bODR = float(params.beta[0]) 

392 

393 paramsOut["mODR2"] = float(params.beta[1]) 

394 paramsOut["bODR2"] = float(params.beta[0]) 

395 

396 paramsOut["mPerp"] = -1.0/paramsOut["mODR2"] 

397 

398 return paramsOut 

399 

400 

401def perpDistance(p1, p2, points): 

402 """Calculate the perpendicular distance to a line from a point 

403 

404 Parameters 

405 ---------- 

406 p1 : `numpy.ndarray` 

407 A point on the line 

408 p2 : `numpy.ndarray` 

409 Another point on the line 

410 points : `zip` 

411 The points to calculate the distance to 

412 

413 Returns 

414 ------- 

415 dists : `list` 

416 The distances from the line to the points. Uses the cross 

417 product to work this out. 

418 """ 

419 dists = [] 

420 for point in points: 

421 point = np.array(point) 

422 distToLine = np.cross(p2 - p1, point - p1)/np.linalg.norm(p2 - p1) 

423 dists.append(distToLine) 

424 

425 return dists 

426 

427 

428def mkColormap(colorNames): 

429 """Make a colormap from the list of color names. 

430 

431 Parameters 

432 ---------- 

433 colorNames : `list` 

434 A list of strings that correspond to matplotlib 

435 named colors. 

436 

437 Returns 

438 ------- 

439 cmap : `matplotlib.colors.LinearSegmentedColormap` 

440 """ 

441 

442 nums = np.linspace(0, 1, len(colorNames)) 

443 blues = [] 

444 greens = [] 

445 reds = [] 

446 for (num, color) in zip(nums, colorNames): 

447 r, g, b = colors.colorConverter.to_rgb(color) 

448 blues.append((num, b, b)) 

449 greens.append((num, g, g)) 

450 reds.append((num, r, r)) 

451 

452 colorDict = {"blue": blues, "red": reds, "green": greens} 

453 cmap = colors.LinearSegmentedColormap("newCmap", colorDict) 

454 return cmap 

455 

456 

457def extremaSort(xs): 

458 """Return the ids of the points reordered so that those 

459 furthest from the median, in absolute terms, are last. 

460 

461 Parameters 

462 ---------- 

463 xs : `np.array` 

464 An array of the values to sort 

465 

466 Returns 

467 ------- 

468 ids : `np.array` 

469 """ 

470 

471 med = np.median(xs) 

472 dists = np.abs(xs - med) 

473 ids = np.argsort(dists) 

474 return ids