Coverage for python/lsst/analysis/ap/plotImageSubtractionCutouts.py: 87%

289 statements  

« prev     ^ index     » next       coverage.py v7.16.0, created at 2026-09-06 09:53 +0000

1# This file is part of analysis_ap. 

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 

22"""Construct template/image/difference cutouts for upload to Zooniverse, or 

23to just to view as images. 

24""" 

25 

26__all__ = ["PlotImageSubtractionCutoutsConfig", "PlotImageSubtractionCutoutsTask", "CutoutPath"] 

27 

28import argparse 

29import functools 

30import io 

31import logging 

32import multiprocessing 

33import os 

34from math import log10 

35 

36import astropy.units as u 

37from lsst.daf.butler import DatasetNotFoundError 

38import lsst.dax.apdb 

39import lsst.geom 

40import lsst.pex.config as pexConfig 

41import lsst.pex.exceptions 

42import lsst.pipe.base 

43import lsst.utils 

44import numpy as np 

45import pandas as pd 

46 

47from . import apdb 

48 

49_log = logging.getLogger(__name__) 

50 

51 

52class _ButlerCache: 

53 """Global class to handle butler queries, to allow lru_cache and 

54 `multiprocessing.Pool` to work together. 

55 

56 If we redo this all to work with BPS or other parallelized systems, or get 

57 good butler-side caching, we could remove this lru_cache system. 

58 """ 

59 

60 def set(self, butler, config): 

61 """Call this to store a Butler and Config instance before using the 

62 global class instance. 

63 

64 Parameters 

65 ---------- 

66 butler : `lsst.daf.butler.Butler` 

67 Butler instance to store. 

68 config : `lsst.pex.config.Config` 

69 Config instance to store. 

70 """ 

71 self._butler = butler 

72 self._config = config 

73 # Ensure the caches are empty if we've been re-set. 

74 self.get_exposures.cache_clear() 

75 self.get_catalog.cache_clear() 

76 

77 @functools.lru_cache(maxsize=4) 

78 def get_exposures(self, instrument, detector, visit): 

79 """Return science, template, difference exposures, using a small 

80 cache so we don't have to re-read files as often. 

81 

82 Parameters 

83 ---------- 

84 instrument : `str` 

85 Instrument name to define the data id. 

86 detector : `int` 

87 Detector id to define the data id. 

88 visit : `int` 

89 Visit id to define the data id. 

90 

91 Returns 

92 ------- 

93 exposures : `tuple` [`lsst.afw.image.ExposureF`] 

94 Science, template, and difference exposure for this data id. 

95 """ 

96 data_id = {'instrument': instrument, 'detector': detector, 'visit': visit} 

97 try: 

98 science = self._butler.get(self._config.science_image_type, data_id) 

99 except DatasetNotFoundError as e: 

100 _log.error(f"Cannot load {self._config.science_image_type} with data_id {data_id}: {e}") 

101 _log.error("If you are working with data processed earlier than May 2025, try setting " 

102 "config.science_image_type = 'initial_pvi' or 'calexp'.") 

103 raise 

104 

105 if self._config.diff_image_type is not None: 105 ↛ 106line 105 didn't jump to line 106 because the condition on line 105 was never true

106 template = self._butler.get(f"{self._config.diff_image_type}_templateExp", data_id) 

107 difference = self._butler.get(f"{self._config.diff_image_type}_differenceExp", data_id) 

108 else: 

109 template = self._butler.get("template_detector", data_id) 

110 difference = self._butler.get("difference_image", data_id) 

111 

112 return science, template, difference 

113 

114 @functools.lru_cache(maxsize=4) 

115 def get_catalog(self, instrument, detector, visit): 

116 """Return the diaSrc catalog from the butler. 

117 

118 Parameters 

119 ---------- 

120 instrument : `str` 

121 Instrument name to define the data id. 

122 detector : `int` 

123 Detector id to define the data id. 

124 visit : `int` 

125 Visit id to define the data id. 

126 

127 Returns 

128 ------- 

129 catalog : `lsst.afw.table.SourceCatalog` 

130 DiaSource catalog for this data id. 

131 """ 

132 data_id = {'instrument': instrument, 'detector': detector, 'visit': visit} 

133 return self._butler.get(f'{self._config.diff_image_type}_diaSrc', data_id) 

134 

135 

136# Global used within each multiprocessing worker (or single process). 

137butler_cache = _ButlerCache() 

138 

139 

140class PlotImageSubtractionCutoutsConfig(pexConfig.Config): 

141 sizes = pexConfig.ListField( 

142 doc="List of widths of cutout to extract for image from science, \ 

143 template, and difference exposures.", 

144 dtype=int, 

145 default=[30], 

146 ) 

147 use_footprint = pexConfig.Field( 

148 doc="Use source footprint to to define cutout region; " 

149 "If set, ignore `size` and use the footprint bbox instead.", 

150 dtype=bool, 

151 default=False, 

152 ) 

153 url_root = pexConfig.Field( 

154 doc="URL that the resulting images will be served to Zooniverse from, for the manifest file. " 

155 "If not set, no manifest file will be written.", 

156 dtype=str, 

157 default=None, 

158 optional=True, 

159 ) 

160 diff_image_type = pexConfig.Field( 

161 doc="Optional partial dataset name of template and difference image to use for cutouts; " 

162 "will have '_templateExp' and '_differenceExp' appended for butler.get(), respectively." 

163 " If not specified, use `template_detector` and `difference_image`, respectively.", 

164 dtype=str, 

165 default=None, 

166 optional=True 

167 ) 

168 science_image_type = pexConfig.Field( 

169 doc="Dataset type of science image to use for cutouts; " 

170 "older processings could be `calexp` or `initial_pvi`.", 

171 dtype=str, 

172 default="preliminary_visit_image", 

173 ) 

174 add_metadata = pexConfig.Field( 

175 doc="Annotate the cutouts with catalog metadata, including coordinates, fluxes, flags, etc.", 

176 dtype=bool, 

177 default=True 

178 ) 

179 chunk_size = pexConfig.Field( 

180 doc="Chunk up files into subdirectories, with at most this many files per directory." 

181 " None means write all the files to one `images/` directory.", 

182 dtype=int, 

183 default=10000, 

184 optional=True 

185 ) 

186 save_as_numpy = pexConfig.Field( 

187 doc="Save the raw cutout images in numpy format.", 

188 dtype=bool, 

189 default=False 

190 ) 

191 

192 

193class PlotImageSubtractionCutoutsTask(lsst.pipe.base.Task): 

194 """Generate template/science/difference image cutouts of DiaSources and an 

195 optional manifest for upload to a Zooniverse project. 

196 

197 Parameters 

198 ---------- 

199 output_path : `str` 

200 The path to write the output to; manifest goes here, while the 

201 images themselves go into ``output_path/images/``. 

202 """ 

203 ConfigClass = PlotImageSubtractionCutoutsConfig 

204 _DefaultName = "plotImageSubtractionCutouts" 

205 

206 def __init__(self, *, output_path, **kwargs): 

207 super().__init__(**kwargs) 

208 self._output_path = output_path 

209 self.cutout_path = CutoutPath(output_path, chunk_size=self.config.chunk_size) 

210 self.numpy_path = CutoutPath(output_path, chunk_size=self.config.chunk_size, 

211 subdirectory='numpy') 

212 

213 def _reduce_kwargs(self): 

214 # to allow pickling of this Task 

215 kwargs = super()._reduce_kwargs() 

216 kwargs["output_path"] = self._output_path 

217 return kwargs 

218 

219 def run(self, data, butler, njobs=0): 

220 """Generate cutout images and a manifest for upload to Zooniverse 

221 from a collection of DiaSources. 

222 

223 Parameters 

224 ---------- 

225 data : `pandas.DataFrame` 

226 The DiaSources to extract cutouts for. Must contain at least these 

227 fields: ``ra, dec, diaSourceId, detector, visit, instrument``. 

228 butler : `lsst.daf.butler.Butler` 

229 The butler connection to use to load the data; create it with the 

230 collections you wish to load images from. 

231 njobs : `int`, optional 

232 Number of multiprocessing jobs to make cutouts with; default of 0 

233 means don't use multiprocessing at all. 

234 

235 Returns 

236 ------- 

237 source_ids : `list` [`int`] 

238 DiaSourceIds of cutout images that were generated. 

239 """ 

240 result = self.write_images(data, butler, njobs=njobs) 

241 self.write_manifest(result) 

242 self.log.info("Wrote %d images to %s", len(result), self._output_path) 

243 return result 

244 

245 def write_manifest(self, sources): 

246 """Save a Zooniverse manifest attaching image URLs to source ids. 

247 

248 Parameters 

249 ---------- 

250 sources : `list` [`int`] 

251 The diaSourceIds of the sources that had cutouts succesfully made. 

252 """ 

253 if self.config.url_root is not None: 

254 manifest = self._make_manifest(sources) 

255 manifest.to_csv(os.path.join(self._output_path, "manifest.csv"), index=False) 

256 else: 

257 self.log.info("No url_root config provided, so no Zooniverse manifest file was written.") 

258 

259 def _make_manifest(self, sources): 

260 """Return a Zooniverse manifest attaching image URLs to source ids. 

261 

262 Parameters 

263 ---------- 

264 sources : `list` [`int`] 

265 The diaSourceIds of the sources that had cutouts succesfully made. 

266 

267 Returns 

268 ------- 

269 manifest : `pandas.DataFrame` 

270 The formatted URL manifest for upload to Zooniverse. 

271 """ 

272 cutout_path = CutoutPath(self.config.url_root) 

273 manifest = pd.DataFrame() 

274 manifest["external_id"] = sources 

275 manifest["location:1"] = [cutout_path(x, f'{x}.png') for x in sources] 

276 manifest["metadata:diaSourceId"] = sources 

277 return manifest 

278 

279 def write_images(self, data, butler, njobs=0): 

280 """Make the 3-part cutout images for each requested source and write 

281 them to disk. 

282 

283 Creates ``images/`` and ``numpy/`` subdirectories if they 

284 do not already exist; images are written there as PNG and npy files. 

285 

286 Parameters 

287 ---------- 

288 data : `pandas.DataFrame` 

289 The DiaSources to extract cutouts for. Must contain at least these 

290 fields: ``ra, dec, diaSourceId, detector, visit, instrument``. 

291 butler : `lsst.daf.butler.Butler` 

292 The butler connection to use to load the data; create it with the 

293 collections you wish to load images from. 

294 njobs : `int`, optional 

295 Number of multiprocessing jobs to make cutouts with; default of 0 

296 means don't use multiprocessing at all. 

297 

298 Returns 

299 ------- 

300 sources : `list` 

301 DiaSourceIds that had cutouts made. 

302 """ 

303 # Ignore divide-by-zero and log-of-negative-value messages. 

304 seterr_dict = np.seterr(divide="ignore", invalid="ignore") 

305 

306 # Exclude index if they are replicated in columns. 

307 indexNotInColumns = not any(index in data.columns for index in data.index.names) 

308 

309 sources = [] 

310 butler_cache.set(butler, self.config) 

311 if njobs > 0: 311 ↛ 312line 311 didn't jump to line 312 because the condition on line 311 was never true

312 with multiprocessing.Pool(njobs) as pool: 

313 sources = pool.map(self._do_one_source, data.to_records(index=indexNotInColumns)) 

314 else: 

315 for source in data.to_records(index=indexNotInColumns): 

316 src_id = source["diaSourceId"] 

317 if self.cutout_path.exists(src_id, f"{src_id}.png"): 317 ↛ 318line 317 didn't jump to line 318 because the condition on line 317 was never true

318 sources.append(src_id) 

319 else: 

320 sources.append(self._do_one_source(source)) 

321 

322 # restore numpy error message state 

323 np.seterr(**seterr_dict) 

324 # Only return successful ids, not failures. 

325 return [s for s in sources if s is not None] 

326 

327 def _do_one_source(self, source): 

328 """Make cutouts for one diaSource. 

329 

330 Parameters 

331 ---------- 

332 source : `numpy.record`, optional 

333 DiaSource record for this cutout, to add metadata to the image. 

334 

335 Returns 

336 ------- 

337 diaSourceId : `int` or None 

338 Id of the source that was generated, or None if there was an error. 

339 """ 

340 try: 

341 center = lsst.geom.SpherePoint(source["ra"], source["dec"], lsst.geom.degrees) 

342 science, template, difference = butler_cache.get_exposures(source["instrument"], 

343 source["detector"], 

344 source["visit"]) 

345 if self.config.use_footprint: 

346 catalog = butler_cache.get_catalog(source["instrument"], 

347 source["detector"], 

348 source["visit"]) 

349 # The input catalogs must be sorted. 

350 if not catalog.isSorted(): 350 ↛ 351line 350 didn't jump to line 351 because the condition on line 350 was never true

351 data_id = {'instrument': source["instrument"], 

352 'detector': source["detector"], 

353 'visit': source["visit"]} 

354 msg = f"{self.config.diff_image_type}_diaSrc catalog for {data_id} is not sorted!" 

355 raise RuntimeError(msg) 

356 record = catalog.find(source['diaSourceId']) 

357 footprint = record.getFootprint() 

358 

359 scale = science.wcs.getPixelScale(science.getBBox().getCenter()).asArcseconds() 

360 image = self.generate_image(science, template, difference, center, scale, 

361 dia_source_id=source['diaSourceId'], 

362 save_as_numpy=self.config.save_as_numpy, 

363 source=source if self.config.add_metadata else None, 

364 footprint=footprint if self.config.use_footprint else None) 

365 self.cutout_path.mkdir(source["diaSourceId"]) 

366 with open(self.cutout_path(source["diaSourceId"], 

367 f'{source["diaSourceId"]}.png'), "wb") as outfile: 

368 outfile.write(image.getbuffer()) 

369 return source["diaSourceId"] 

370 except (LookupError, lsst.pex.exceptions.Exception) as e: 

371 self.log.error( 

372 f"{e.__class__.__name__} processing diaSourceId {source['diaSourceId']}: {e}" 

373 ) 

374 return None 

375 except Exception: 

376 # Ensure other exceptions are interpretable when multiprocessing. 

377 import traceback 

378 traceback.print_exc() 

379 raise 

380 

381 def generate_image(self, science, template, difference, center, scale, dia_source_id=None, 

382 save_as_numpy=False, source=None, footprint=None): 

383 """Get a 3-part cutout image to save to disk, for a single source. 

384 

385 Parameters 

386 ---------- 

387 science : `lsst.afw.image.ExposureF` 

388 Science exposure to include in the cutout. 

389 template : `lsst.afw.image.ExposureF` 

390 Matched template exposure to include in the cutout. 

391 difference : `lsst.afw.image.ExposureF` 

392 Matched science minus template exposure to include in the cutout. 

393 center : `lsst.geom.SpherePoint` 

394 Center of the source to be cut out of each image. 

395 scale : `float` 

396 Pixel scale in arcseconds. 

397 dia_source_id : `int`, optional 

398 DiaSourceId to use in the filename, if saving to disk. 

399 save_as_numpy : `bool`, optional 

400 Save the raw cutout images in numpy format. 

401 source : `numpy.record`, optional 

402 DiaSource record for this cutout, to add metadata to the image. 

403 footprint : `lsst.afw.detection.Footprint`, optional 

404 Detected source footprint; if specified, extract a square 

405 surrounding the footprint bbox, otherwise use ``config.size``. 

406 

407 Returns 

408 ------- 

409 image: `io.BytesIO` 

410 The generated image, to be output to a file or displayed on screen. 

411 """ 

412 numpy_cutouts = {} 

413 if not self.config.use_footprint: 

414 sizes = self.config.sizes 

415 cutout_science, cutout_template, cutout_difference = [], [], [] 

416 for i, s in enumerate(sizes): 

417 extent = lsst.geom.Extent2I(s, s) 

418 science_cutout = science.getCutout(center, extent) 

419 template_cutout = template.getCutout(center, extent) 

420 difference_cutout = difference.getCutout(center, extent) 

421 if save_as_numpy: 

422 self.numpy_path.mkdir(dia_source_id) 

423 numpy_cutouts[f"sci_{s}"] = science_cutout.image.array 

424 numpy_cutouts[f"temp_{s}"] = template_cutout.image.array 

425 numpy_cutouts[f"diff_{s}"] = difference_cutout.image.array 

426 for cutout_type, cutout in numpy_cutouts.items(): 

427 outfile = self.numpy_path(dia_source_id, f'{dia_source_id}_{cutout_type}.npy') 

428 np.save(outfile, np.expand_dims(cutout, axis=0)) 

429 cutout_science.append(science_cutout) 

430 cutout_template.append(template_cutout) 

431 cutout_difference.append(difference_cutout) 

432 else: 

433 if self.config.save_as_numpy: 433 ↛ 434line 433 didn't jump to line 434 because the condition on line 433 was never true

434 raise RuntimeError("Cannot save as numpy when using footprints.") 

435 cutout_science = [science.getCutout(footprint.getBBox())] 

436 cutout_template = [template.getCutout(footprint.getBBox())] 

437 cutout_difference = [difference.getCutout(footprint.getBBox())] 

438 extent = footprint.getBBox().getDimensions() 

439 # Plot a square equal to the largest dimension. 

440 sizes = [extent.x if extent.x > extent.y else extent.y] 

441 

442 return self._plot_cutout(cutout_science, 

443 cutout_template, 

444 cutout_difference, 

445 scale, 

446 sizes, 

447 source=source) 

448 

449 def _plot_cutout(self, science, template, difference, scale, sizes, source=None): 

450 """Plot the cutouts for a source in one image. 

451 

452 Parameters 

453 ---------- 

454 science : `list` [`lsst.afw.image.ExposureF`] 

455 List of cutout Science exposure(s) to include in the image. 

456 template : `list` [`lsst.afw.image.ExposureF`] 

457 List of cutout template exposure(s) to include in the image. 

458 difference : `list` [`lsst.afw.image.ExposureF`] 

459 List of cutout science minus template exposure(s) to include 

460 in the image. 

461 source : `numpy.record`, optional 

462 DiaSource record for this cutout, to add metadata to the image. 

463 scale : `float` 

464 Pixel scale in arcseconds. 

465 size : `list` [`int`] 

466 List of x/y dimensions of of the images passed in, to set imshow 

467 extent. 

468 

469 Returns 

470 ------- 

471 image: `io.BytesIO` 

472 The generated image, to be output to a file via 

473 `image.write(filename)` or displayed on screen. 

474 """ 

475 import astropy.visualization as aviz 

476 import matplotlib 

477 matplotlib.use("AGG") 

478 # Force matplotlib defaults 

479 matplotlib.rcParams.update(matplotlib.rcParamsDefault) 

480 import matplotlib.pyplot as plt 

481 from matplotlib import cm 

482 

483 # TODO DM-32014: how do we color masked pixels (including edges)? 

484 

485 def plot_one_image(ax, data, size, name=None): 

486 """Plot a normalized image on an axis.""" 

487 if name == "Difference": 

488 norm = aviz.ImageNormalize( 

489 # focus on a rect of dim 15 at the center of the image. 

490 data[data.shape[0] // 2 - 7:data.shape[0] // 2 + 8, 

491 data.shape[1] // 2 - 7:data.shape[1] // 2 + 8], 

492 interval=aviz.MinMaxInterval(), 

493 stretch=aviz.AsinhStretch(a=0.1), 

494 ) 

495 else: 

496 norm = aviz.ImageNormalize( 

497 data, 

498 interval=aviz.MinMaxInterval(), 

499 stretch=aviz.AsinhStretch(a=0.1), 

500 ) 

501 ax.imshow(data, cmap=cm.bone, interpolation="none", norm=norm, 

502 extent=(0, size, 0, size), origin="lower", aspect="equal") 

503 x_line = 1 

504 y_line = 1 

505 ax.plot((x_line, x_line + 1.0/scale), (y_line, y_line), color="blue", lw=6) 

506 ax.plot((x_line, x_line + 1.0/scale), (y_line, y_line), color="yellow", lw=2) 

507 ax.axis("off") 

508 if name is not None: 

509 ax.set_title(name) 

510 

511 try: 

512 len_sizes = len(sizes) 

513 fig, axs = plt.subplots(len_sizes, 3, constrained_layout=True) 

514 if len_sizes == 1: 

515 plot_one_image(axs[0], template[0].image.array, sizes[0], "Template") 

516 plot_one_image(axs[1], science[0].image.array, sizes[0], "Science") 

517 plot_one_image(axs[2], difference[0].image.array, sizes[0], "Difference") 

518 else: 

519 plot_one_image(axs[0][0], template[0].image.array, sizes[0], "Template") 

520 plot_one_image(axs[0][1], science[0].image.array, sizes[0], "Science") 

521 plot_one_image(axs[0][2], difference[0].image.array, sizes[0], "Difference") 

522 for i in range(1, len(axs)): 

523 plot_one_image(axs[i][0], template[i].image.array, sizes[i], None) 

524 plot_one_image(axs[i][1], science[i].image.array, sizes[i], None) 

525 plot_one_image(axs[i][2], difference[i].image.array, sizes[i], None) 

526 if source is not None: 

527 _annotate_image(fig, source, len_sizes) 

528 

529 output = io.BytesIO() 

530 plt.savefig(output, bbox_inches="tight", format="png") 

531 output.seek(0) # to ensure opening the image starts from the front 

532 finally: 

533 plt.close(fig) 

534 

535 return output 

536 

537 

538# Flag groupings for the metadata legend on cutout images. A row's label is 

539# colored red if any column in its group is set on the source. 

540_FLAG_GROUPS = { 

541 "psf": ["psfFlux_flag", "psfFlux_flag_noGoodPixels", "psfFlux_flag_edge"], 

542 "aperture": ["apFlux_flag", "apFlux_flag_apertureTruncated"], 

543 "forced": ["forced_PsfFlux_flag", "forced_PsfFlux_flag_noGoodPixels", 

544 "forced_PsfFlux_flag_edge"], 

545 "edge": ["pixelFlags_edge"], 

546 "interp": ["pixelFlags_interpolated", "pixelFlags_interpolatedCenter"], 

547 "saturated": ["pixelFlags_saturated", "pixelFlags_saturatedCenter"], 

548 "cr": ["pixelFlags_cr", "pixelFlags_crCenter"], 

549 "bad": ["pixelFlags_bad"], 

550 "suspect": ["pixelFlags_suspect", "pixelFlags_suspectCenter"], 

551 "centroid": ["centroid_flag"], 

552 "shape": ["shape_flag", "shape_flag_no_pixels", "shape_flag_not_contained", 

553 "shape_flag_parent_source"], 

554} 

555 

556# Flag-tag overlays drawn on top of the flux rows (rows 2 and 3). Each entry 

557# is (predicate, x, label, color, row_index). The predicate is either a key 

558# into ``_FLAG_GROUPS`` (the tag is drawn if any column in the group is set) 

559# or a callable returning a bool. 

560_FLAG_TAGS = [ 

561 ("edge", 0.55, "EDGE", "goldenrod", 2), 

562 ("interp", 0.62, "INTERP", "green", 2), 

563 ("saturated", 0.72, "SAT", "green", 2), 

564 ("cr", 0.77, "CR", "magenta", 2), 

565 ("bad", 0.81, "BAD", "red", 2), 

566 (lambda src: bool(src["isDipole"]), 0.87, "DIPOLE", "indigo", 2), 

567 ("suspect", 0.55, "SUS", "goldenrod", 3), 

568 ("centroid", 0.60, "CENTROID", "red", 3), 

569 ("shape", 0.73, "SHAPE", "red", 3), 

570] 

571# Future option: add two more flag flavors at x = 0.80 and 0.87 on row 3. 

572 

573 

574def _annotate_image(fig, source, len_sizes, heights=None): 

575 """Annotate the cutouts image with metadata and flags. 

576 

577 Parameters 

578 ---------- 

579 fig : `matplotlib.Figure` 

580 Figure to be annotated. 

581 source : `numpy.record` 

582 DiaSource record of the object being plotted. 

583 len_sizes : `int` 

584 Length of the ``size`` array set in configuration. 

585 heights : `list` [`float`], optional 

586 Five figure-fraction y-coordinates for the metadata rows. If None, 

587 the default heights are chosen based on ``len_sizes``. Subclasses 

588 that add extra panels to the figure can pass their own positions. 

589 """ 

590 flag_color = "red" 

591 text_color = "grey" 

592 

593 if heights is None: 

594 if len_sizes == 1: 

595 heights = [0.95, 0.91, 0.87, 0.83, 0.79] 

596 else: 

597 heights = [1.2, 1.15, 1.1, 1.05, 1.0] 

598 

599 def label_color(group_key): 

600 """Red label if any flag in the group is set, otherwise grey.""" 

601 return flag_color if any(source[_FLAG_GROUPS[group_key]]) else text_color 

602 

603 # Each row is a list of (x, text, kwargs) atoms drawn at heights[row_idx]. 

604 # fig.text coordinates are in fractions of the figure. 

605 rows = [ 

606 # Row 0: identity (diaSourceId, instrument, detector, visit, band). 

607 [ 

608 (0.000, "diaSourceId:", {"color": text_color}), 

609 (0.145, f"{source['diaSourceId']}", {}), 

610 (0.430, f"{source['instrument']}", {"fontweight": "bold"}), 

611 (0.640, "detector:", {"color": text_color}), 

612 (0.740, f"{source['detector']}", {}), 

613 (0.795, "visit:", {"color": text_color}), 

614 (0.850, f"{source['visit']}", {}), 

615 (0.950, f"{source['band']}", {}), 

616 ], 

617 # Row 1: coordinates and detection-quality numbers. 

618 [ 

619 (0.000, "ra:", {"color": text_color}), 

620 (0.037, f"{source['ra']:.8f}", {}), 

621 (0.210, "dec:", {"color": text_color}), 

622 (0.265, f"{source['dec']:+.8f}", {}), 

623 (0.500, "detection S/N:", {"color": text_color}), 

624 (0.660, f"{source['snr']:6.1f}", {}), 

625 (0.750, "PSF chi2:", {"color": text_color}), 

626 (0.850, f"{source['psfChi2']/source['psfNdata']:6.2f}", {}), 

627 ], 

628 # Row 2: PSF flux. 

629 [ 

630 (0.000, "PSF (nJy):", {"color": label_color("psf")}), 

631 (0.250, f"{source['psfFlux']:8.1f}", {"horizontalalignment": "right"}), 

632 (0.252, "+/-", {"color": text_color}), 

633 (0.290, f"{source['psfFluxErr']:8.1f}", {}), 

634 (0.400, "S/N:", {"color": text_color}), 

635 (0.450, f"{abs(source['psfFlux']/source['psfFluxErr']):6.2f}", {}), 

636 ], 

637 # Row 3: aperture flux. 

638 [ 

639 (0.000, "ap (nJy):", {"color": label_color("aperture")}), 

640 (0.250, f"{source['apFlux']:8.1f}", {"horizontalalignment": "right"}), 

641 (0.252, "+/-", {"color": text_color}), 

642 (0.290, f"{source['apFluxErr']:8.1f}", {}), 

643 (0.400, "S/N:", {"color": text_color}), 

644 (0.450, f"{abs(source['apFlux']/source['apFluxErr']):#6.2f}", {}), 

645 ], 

646 # Row 4: forced-photometry flux + ABmag. 

647 [ 

648 (0.000, "sci (nJy):", {"color": label_color("forced")}), 

649 (0.250, f"{source['scienceFlux']:8.1f}", {"horizontalalignment": "right"}), 

650 (0.252, "+/-", {"color": text_color}), 

651 (0.290, f"{source['scienceFluxErr']:8.1f}", {}), 

652 (0.400, "S/N:", {"color": text_color}), 

653 (0.450, f"{abs(source['scienceFlux']/source['scienceFluxErr']):6.2f}", {}), 

654 (0.550, "ABmag:", {"color": text_color}), 

655 (0.635, f"{(source['scienceFlux']*u.nanojansky).to_value(u.ABmag):.3f}", {}), 

656 ], 

657 ] 

658 for row, y in zip(rows, heights): 

659 for x, text, kwargs in row: 

660 fig.text(x, y, text, **kwargs) 

661 

662 # Draw flag-tag overlays after the rows so they sit on top. 

663 for predicate, x, text, color, row_idx in _FLAG_TAGS: 

664 if callable(predicate): 

665 triggered = predicate(source) 

666 else: 

667 triggered = any(source[_FLAG_GROUPS[predicate]]) 

668 if triggered: 

669 fig.text(x, heights[row_idx], text, color=color, fontweight="bold") 

670 

671 # Reliability score: color depends on the value, not on flags. 

672 if source['reliability'] is not None and np.isfinite(source['reliability']): 672 ↛ exitline 672 didn't return from function '_annotate_image' because the condition on line 672 was always true

673 rb = source['reliability'] 

674 fig.text(0.73, heights[4], f"RB:{rb:.03f}", 

675 color='#e41a1c' if rb < 0.5 else '#4daf4a', 

676 fontweight="bold") 

677 

678 

679class CutoutPath: 

680 """Manage paths to image cutouts with filenames based on diaSourceId. 

681 

682 Supports local files, and id-chunked directories. 

683 

684 Parameters 

685 ---------- 

686 root : `str` 

687 Root file path to manage. 

688 chunk_size : `int`, optional 

689 At most this many files per directory. Must be a power of 10. 

690 subdirectory : `str`, optional 

691 Name of the subdirectory 

692 

693 Raises 

694 ------ 

695 RuntimeError 

696 Raised if chunk_size is not a power of 10. 

697 """ 

698 

699 def __init__(self, root, chunk_size=None, subdirectory='images'): 

700 self._root = root 

701 if chunk_size is not None and (log10(chunk_size) != int(log10(chunk_size))): 

702 raise RuntimeError(f"CutoutPath file chunk_size must be a power of 10, got {chunk_size}.") 

703 self._chunk_size = chunk_size 

704 self._subdirectory = subdirectory 

705 

706 def directory(self, id): 

707 """Return the directory to store the output in. 

708 

709 Parameters 

710 ---------- 

711 id : `int` 

712 Source id to create the path for. 

713 

714 Returns 

715 ------- 

716 directory: `str` 

717 Directory for this file. 

718 """ 

719 

720 def chunker(id, size): 

721 return (id // size)*size 

722 

723 if self._chunk_size is not None: 

724 return os.path.join(self._root, 

725 f"{self._subdirectory}/{chunker(id, self._chunk_size)}") 

726 else: 

727 return os.path.join(self._root, f"{self._subdirectory}") 

728 

729 def __call__(self, id, filename): 

730 """Return the full path to a diaSource cutout. 

731 

732 Parameters 

733 ---------- 

734 id : `int` 

735 Source id to create the path for. 

736 filename: `str` 

737 Filename to write. 

738 

739 Returns 

740 ------- 

741 path : `str` 

742 Full path to the requested file. 

743 """ 

744 

745 return os.path.join(self.directory(id), filename) 

746 

747 def exists(self, id, filename): 

748 """Return True if the file already exists. 

749 

750 Parameters 

751 ---------- 

752 id : `int` 

753 Source id to create the path for. 

754 filename: `str` 

755 Filename to write. 

756 

757 Returns 

758 ------- 

759 exists : `bool` 

760 Does the supplied filename exist? 

761 """ 

762 

763 return os.path.exists(os.path.join(self.directory(id), filename)) 

764 

765 def mkdir(self, id): 

766 """Make the directory tree to write this cutout id to. 

767 

768 Parameters 

769 ---------- 

770 id : `int` 

771 Source id to create the path for. 

772 """ 

773 os.makedirs(self.directory(id), exist_ok=True) 

774 

775 

776def build_argparser(): 

777 """Construct an argument parser for the ``plotImageSubtractionCutouts`` 

778 script. 

779 

780 Returns 

781 ------- 

782 argparser : `argparse.ArgumentParser` 

783 The argument parser that defines the ``plotImageSubtractionCutouts`` 

784 command-line interface. 

785 """ 

786 parser = argparse.ArgumentParser( 

787 description=__doc__, 

788 formatter_class=argparse.RawDescriptionHelpFormatter, 

789 epilog="More information is available at https://pipelines.lsst.io.", 

790 ) 

791 

792 apdbArgs = parser.add_mutually_exclusive_group(required=True) 

793 apdbArgs.add_argument( 

794 "--sqlitefile", 

795 default=None, 

796 help="Path to sqlite file to load from; required for sqlite connection.", 

797 ) 

798 apdbArgs.add_argument( 

799 "--namespace", 

800 default=None, 

801 help="Postgres namespace (aka schema) to connect to; " 

802 " required for postgres connections." 

803 ) 

804 

805 parser.add_argument( 

806 "--postgres_url", 

807 default="rubin@usdf-prompt-processing-dev.slac.stanford.edu/lsst-devl", 

808 help="Postgres connection path, or default (None) to use ApdbPostgresQuery default." 

809 ) 

810 

811 parser.add_argument( 

812 "--limit", 

813 default=5, 

814 type=int, 

815 help="Number of sources to load from the APDB (default=5), or the " 

816 "number of sources to load per 'page' when `--all` is set. " 

817 "This should be significantly larger (100x or more) than the value of `-j`, " 

818 "to ensure efficient use of each process.", 

819 ) 

820 parser.add_argument( 

821 "--all", 

822 default=False, 

823 action="store_true", 

824 help="Process all the sources; --limit then becomes the 'page size' to chunk the DB into.", 

825 ) 

826 

827 parser.add_argument( 

828 "-j", 

829 "--jobs", 

830 default=0, 

831 type=int, 

832 help="Number of processes to use when generating cutouts. " 

833 "Specify 0 (the default) to not use multiprocessing at all. " 

834 "Note that `--limit` determines how efficiently each process is filled." 

835 ) 

836 

837 parser.add_argument( 

838 "-C", 

839 "--configFile", 

840 help="File containing the PlotImageSubtractionCutoutsConfig to load.", 

841 ) 

842 parser.add_argument( 

843 "--collections", 

844 nargs="*", 

845 help=( 

846 "Butler collection(s) to load data from." 

847 " If not specified, will search all butler collections, " 

848 "which may be very slow." 

849 ), 

850 ) 

851 parser.add_argument("repo", help="Path to Butler repository to load data from.") 

852 parser.add_argument( 

853 "outputPath", 

854 help="Path to write the output images and manifest to; " 

855 "manifest is written here, while the images go to `OUTPUTPATH/images/`.", 

856 ) 

857 parser.add_argument( 

858 "--reliabilityMin", 

859 type=float, 

860 default=None, 

861 help="Minimum reliability value (default=None) on which to filter the DiaSources.", 

862 ) 

863 parser.add_argument( 

864 "--reliabilityMax", 

865 type=float, 

866 default=None, 

867 help="Maximum reliability value (default=None) on which to filter the DiaSources.", 

868 ) 

869 return parser 

870 

871 

872def _make_apdbQuery(sqlitefile=None, postgres_url=None, namespace=None): 

873 """Return a query connection to the specified APDB. 

874 

875 Parameters 

876 ---------- 

877 sqlitefile : `str`, optional 

878 SQLite file to load APDB from; if set, postgres kwargs are ignored. 

879 postgres_url : `str`, optional 

880 Postgres connection URL to connect to APDB. 

881 namespace : `str`, optional 

882 Postgres schema to load from; required with postgres_url. 

883 

884 Returns 

885 ------- 

886 apdb_query : `lsst.analysis.ap.ApdbQuery` 

887 Query instance to use to load data from APDB. 

888 

889 Raises 

890 ------ 

891 RuntimeError 

892 Raised if the APDB connection kwargs are invalid in some way. 

893 """ 

894 if sqlitefile is not None: 894 ↛ 896line 894 didn't jump to line 896 because the condition on line 894 was always true

895 apdb_query = apdb.ApdbSqliteQuery(sqlitefile) 

896 elif postgres_url is not None and namespace is not None: 

897 apdb_query = apdb.ApdbPostgresQuery(namespace, postgres_url) 

898 else: 

899 raise RuntimeError("Cannot handle database connection args: " 

900 f"sqlitefile={sqlitefile}, postgres_url={postgres_url}, namespace={namespace}") 

901 return apdb_query 

902 

903 

904def run_cutouts(args): 

905 """Run PlotImageSubtractionCutoutsTask on the parsed commandline arguments. 

906 

907 Parameters 

908 ---------- 

909 args : `argparse.Namespace` 

910 Parsed commandline arguments. 

911 """ 

912 # We have to initialize the logger manually on the commandline. 

913 logging.basicConfig( 

914 level=logging.INFO, format="{name} {levelname}: {message}", style="{" 

915 ) 

916 

917 butler = lsst.daf.butler.Butler(args.repo, collections=args.collections) 

918 apdb_query = _make_apdbQuery(sqlitefile=args.sqlitefile, 

919 postgres_url=args.postgres_url, 

920 namespace=args.namespace) 

921 

922 config = PlotImageSubtractionCutoutsConfig() 

923 if args.configFile is not None: 923 ↛ 925line 923 didn't jump to line 925 because the condition on line 923 was always true

924 config.load(os.path.expanduser(args.configFile)) 

925 config.freeze() 

926 cutouts = PlotImageSubtractionCutoutsTask(config=config, output_path=args.outputPath) 

927 

928 if config.save_as_numpy: 928 ↛ 930line 928 didn't jump to line 930 because the condition on line 928 was never true

929 # save the RB output up front so we can use partial runs 

930 data = apdb_query.iter_sources(args.limit, args.reliabilityMin, args.reliabilityMax) 

931 cols_to_export = ["diaSourceId", "visit", "detector", "diaObjectId", 

932 "ssObjectId", "midpointMjdTai", "ra", "dec", "x", "y", 

933 "apFlux", "apFluxErr", "snr", "psfFlux", "psfFluxErr", 

934 "isDipole", "trailLength", "band", "extendedness", 

935 "pixelFlags_bad", "pixelFlags_cr", "pixelFlags_crCenter", 

936 "pixelFlags_edge", "pixelFlags_interpolated", "pixelFlags_interpolatedCenter", 

937 "pixelFlags_offimage", "pixelFlags_saturated", "pixelFlags_saturatedCenter", 

938 "pixelFlags_suspect", "pixelFlags_suspectCenter", "pixelFlags_streak", 

939 "pixelFlags_streakCenter", "pixelFlags_injected", "pixelFlags_injectedCenter", 

940 "pixelFlags_injected_template", "pixelFlags_injected_templateCenter"] 

941 # this is inefficient but otherwise we don't use the same query 

942 all_data = pd.concat([d[cols_to_export] for d in data]) 

943 all_data.to_csv(os.path.join(args.outputPath, "all_diasources.csv.gz"), index=False) 

944 

945 getter = apdb_query.iter_sources(args.limit, args.reliabilityMin, args.reliabilityMax) 

946 # Process just one block of length "limit", or all sources in the database? 

947 if not args.all: 

948 data = next(getter) 

949 sources = cutouts.run(data, butler, njobs=args.jobs) 

950 else: 

951 sources = [] 

952 count = apdb_query.count_sources() 

953 for i, data in enumerate(getter): 

954 sources.extend(cutouts.write_images(data, butler, njobs=args.jobs)) 

955 print(f"Completed {i+1} batches of {args.limit} size, out of {count} diaSources.") 

956 cutouts.write_manifest(sources) 

957 

958 if config.save_as_numpy: 958 ↛ 960line 958 didn't jump to line 960 because the condition on line 958 was never true

959 # Write a dataframe with only diasources successfully written. 

960 data.loc[data['diaSourceId'].isin(sources), cols_to_export].to_csv( 

961 os.path.join(args.outputPath, "exported_diasources.csv.gz"), index=False) 

962 

963 print(f"Generated {len(sources)} diaSource cutouts to {args.outputPath}.") 

964 

965 

966def main(): 

967 args = build_argparser().parse_args() 

968 run_cutouts(args)