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

296 statements  

« prev     ^ index     » next       coverage.py v7.16.2, created at 2026-10-23 09:35 +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 

44from lsst.images import Box, NoOverlapError 

45import numpy as np 

46import pandas as pd 

47 

48from . import apdb 

49from .imageUtils import as_visit_image, cutout, cutout_box, get_visit_image, pixel_scale 

50 

51_log = logging.getLogger(__name__) 

52 

53 

54class _ButlerCache: 

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

56 `multiprocessing.Pool` to work together. 

57 

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

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

60 """ 

61 

62 def set(self, butler, config): 

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

64 global class instance. 

65 

66 Parameters 

67 ---------- 

68 butler : `lsst.daf.butler.Butler` 

69 Butler instance to store. 

70 config : `lsst.pex.config.Config` 

71 Config instance to store. 

72 """ 

73 self._butler = butler 

74 self._config = config 

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

76 self.get_images.cache_clear() 

77 self.get_catalog.cache_clear() 

78 

79 @functools.lru_cache(maxsize=4) 

80 def get_images(self, instrument, detector, visit): 

81 """Return science, template, difference images, using a small 

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

83 

84 Parameters 

85 ---------- 

86 instrument : `str` 

87 Instrument name to define the data id. 

88 detector : `int` 

89 Detector id to define the data id. 

90 visit : `int` 

91 Visit id to define the data id. 

92 

93 Returns 

94 ------- 

95 exposures : `tuple` [`lsst.images.VisitImage`] 

96 Science, template, and difference image for this data id, with 

97 legacy exposures converted by `get_visit_image`. 

98 """ 

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

100 try: 

101 science = get_visit_image(self._butler, self._config.science_image_type, data_id) 

102 except DatasetNotFoundError as e: 

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

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

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

106 raise 

107 

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

109 template_name = f"{self._config.diff_image_type}_templateExp" 

110 difference_name = f"{self._config.diff_image_type}_differenceExp" 

111 else: 

112 template_name = "template_detector" 

113 difference_name = "difference_image" 

114 template = get_visit_image(self._butler, template_name, data_id, image_type="template", 

115 detector=science.detector) 

116 difference = get_visit_image(self._butler, difference_name, data_id, image_type="difference") 

117 

118 return science, template, difference 

119 

120 @functools.lru_cache(maxsize=4) 

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

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

123 

124 Parameters 

125 ---------- 

126 instrument : `str` 

127 Instrument name to define the data id. 

128 detector : `int` 

129 Detector id to define the data id. 

130 visit : `int` 

131 Visit id to define the data id. 

132 

133 Returns 

134 ------- 

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

136 DiaSource catalog for this data id. 

137 """ 

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

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

140 

141 

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

143butler_cache = _ButlerCache() 

144 

145 

146class PlotImageSubtractionCutoutsConfig(pexConfig.Config): 

147 sizes = pexConfig.ListField( 

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

149 template, and difference exposures.", 

150 dtype=int, 

151 default=[30], 

152 ) 

153 use_footprint = pexConfig.Field( 

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

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

156 dtype=bool, 

157 default=False, 

158 ) 

159 url_root = pexConfig.Field( 

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

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

162 dtype=str, 

163 default=None, 

164 optional=True, 

165 ) 

166 diff_image_type = pexConfig.Field( 

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

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

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

170 dtype=str, 

171 default=None, 

172 optional=True 

173 ) 

174 science_image_type = pexConfig.Field( 

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

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

177 dtype=str, 

178 default="preliminary_visit_image", 

179 ) 

180 add_metadata = pexConfig.Field( 

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

182 dtype=bool, 

183 default=True 

184 ) 

185 chunk_size = pexConfig.Field( 

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

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

188 dtype=int, 

189 default=10000, 

190 optional=True 

191 ) 

192 save_as_numpy = pexConfig.Field( 

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

194 dtype=bool, 

195 default=False 

196 ) 

197 

198 

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

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

201 optional manifest for upload to a Zooniverse project. 

202 

203 Parameters 

204 ---------- 

205 output_path : `str` 

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

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

208 """ 

209 ConfigClass = PlotImageSubtractionCutoutsConfig 

210 _DefaultName = "plotImageSubtractionCutouts" 

211 

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

213 super().__init__(**kwargs) 

214 self._output_path = output_path 

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

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

217 subdirectory='numpy') 

218 

219 def _reduce_kwargs(self): 

220 # to allow pickling of this Task 

221 kwargs = super()._reduce_kwargs() 

222 kwargs["output_path"] = self._output_path 

223 return kwargs 

224 

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

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

227 from a collection of DiaSources. 

228 

229 Parameters 

230 ---------- 

231 data : `pandas.DataFrame` 

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

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

234 butler : `lsst.daf.butler.Butler` 

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

236 collections you wish to load images from. 

237 njobs : `int`, optional 

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

239 means don't use multiprocessing at all. 

240 

241 Returns 

242 ------- 

243 source_ids : `list` [`int`] 

244 DiaSourceIds of cutout images that were generated. 

245 """ 

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

247 self.write_manifest(result) 

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

249 return result 

250 

251 def write_manifest(self, sources): 

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

253 

254 Parameters 

255 ---------- 

256 sources : `list` [`int`] 

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

258 """ 

259 if self.config.url_root is not None: 

260 manifest = self._make_manifest(sources) 

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

262 else: 

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

264 

265 def _make_manifest(self, sources): 

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

267 

268 Parameters 

269 ---------- 

270 sources : `list` [`int`] 

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

272 

273 Returns 

274 ------- 

275 manifest : `pandas.DataFrame` 

276 The formatted URL manifest for upload to Zooniverse. 

277 """ 

278 cutout_path = CutoutPath(self.config.url_root) 

279 manifest = pd.DataFrame() 

280 manifest["external_id"] = sources 

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

282 manifest["metadata:diaSourceId"] = sources 

283 return manifest 

284 

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

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

287 them to disk. 

288 

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

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

291 

292 Parameters 

293 ---------- 

294 data : `pandas.DataFrame` 

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

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

297 butler : `lsst.daf.butler.Butler` 

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

299 collections you wish to load images from. 

300 njobs : `int`, optional 

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

302 means don't use multiprocessing at all. 

303 

304 Returns 

305 ------- 

306 sources : `list` 

307 DiaSourceIds that had cutouts made. 

308 """ 

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

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

311 

312 # Exclude index if they are replicated in columns. 

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

314 

315 sources = [] 

316 butler_cache.set(butler, self.config) 

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

318 with multiprocessing.Pool(njobs) as pool: 

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

320 else: 

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

322 src_id = source["diaSourceId"] 

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

324 sources.append(src_id) 

325 else: 

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

327 

328 # restore numpy error message state 

329 np.seterr(**seterr_dict) 

330 # Only return successful ids, not failures. 

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

332 

333 def _do_one_source(self, source): 

334 """Make cutouts for one diaSource. 

335 

336 Parameters 

337 ---------- 

338 source : `numpy.record`, optional 

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

340 

341 Returns 

342 ------- 

343 diaSourceId : `int` or None 

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

345 """ 

346 try: 

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

348 science, template, difference = butler_cache.get_images(source["instrument"], 

349 source["detector"], 

350 source["visit"]) 

351 if self.config.use_footprint: 

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

353 source["detector"], 

354 source["visit"]) 

355 # The input catalogs must be sorted. 

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

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

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

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

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

361 raise RuntimeError(msg) 

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

363 footprint = record.getFootprint() 

364 

365 scale = pixel_scale(science) 

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

367 dia_source_id=source['diaSourceId'], 

368 save_as_numpy=self.config.save_as_numpy, 

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

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

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

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

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

374 outfile.write(image.getbuffer()) 

375 return source["diaSourceId"] 

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

377 self.log.error( 

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

379 ) 

380 return None 

381 except Exception: 

382 # Ensure other exceptions are interpretable when multiprocessing. 

383 import traceback 

384 traceback.print_exc() 

385 raise 

386 

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

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

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

390 

391 Parameters 

392 ---------- 

393 science : `lsst.images.VisitImage` 

394 Science image to include in the cutout. 

395 template : `lsst.images.DifferenceImage` 

396 Matched template image to include in the cutout. 

397 difference : `lsst.images.DifferenceImage` 

398 Matched science minus template image to include in the cutout. 

399 center : `astropy.coordinates.SkyCoord` or `lsst.geom.SpherePoint` 

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

401 scale : `float` 

402 Pixel scale in arcseconds. 

403 dia_source_id : `int`, optional 

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

405 save_as_numpy : `bool`, optional 

406 Save the raw cutout images in numpy format. 

407 source : `numpy.record`, optional 

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

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

410 Detected source footprint; if specified, extract a square 

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

412 

413 Returns 

414 ------- 

415 image: `io.BytesIO` 

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

417 """ 

418 science = as_visit_image(science) 

419 template = as_visit_image(template, image_type="template", detector=science.detector) 

420 difference = as_visit_image(difference, image_type="difference") 

421 numpy_cutouts = {} 

422 if not self.config.use_footprint: 

423 sizes = self.config.sizes 

424 cutout_science, cutout_template, cutout_difference = [], [], [] 

425 for i, s in enumerate(sizes): 

426 science_cutout = cutout(science, center, s) 

427 template_cutout = cutout(template, center, s) 

428 difference_cutout = cutout(difference, center, s) 

429 if save_as_numpy: 

430 self.numpy_path.mkdir(dia_source_id) 

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

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

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

434 for cutout_type, array in numpy_cutouts.items(): 

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

436 np.save(outfile, np.expand_dims(array, axis=0)) 

437 cutout_science.append(science_cutout) 

438 cutout_template.append(template_cutout) 

439 cutout_difference.append(difference_cutout) 

440 else: 

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

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

443 box = Box.from_legacy(footprint.getBBox()) 

444 cutout_science = [cutout_box(science, box)] 

445 cutout_template = [cutout_box(template, box)] 

446 cutout_difference = [cutout_box(difference, box)] 

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

448 # Plot a square equal to the largest dimension. 

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

450 

451 return self._plot_cutout(cutout_science, 

452 cutout_template, 

453 cutout_difference, 

454 scale, 

455 sizes, 

456 source=source) 

457 

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

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

460 

461 Parameters 

462 ---------- 

463 science : `list` [`lsst.images.Image`] 

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

465 template : `list` [`lsst.images.Image`] 

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

467 difference : `list` [`lsst.images.Image`] 

468 List of cutout science minus template image(s) to include 

469 in the image. 

470 source : `numpy.record`, optional 

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

472 scale : `float` 

473 Pixel scale in arcseconds. 

474 size : `list` [`int`] 

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

476 extent. 

477 

478 Returns 

479 ------- 

480 image: `io.BytesIO` 

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

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

483 """ 

484 import astropy.visualization as aviz 

485 import matplotlib 

486 matplotlib.use("AGG") 

487 # Force matplotlib defaults 

488 matplotlib.rcParams.update(matplotlib.rcParamsDefault) 

489 import matplotlib.pyplot as plt 

490 from matplotlib import cm 

491 

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

493 

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

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

496 if name == "Difference": 

497 norm = aviz.ImageNormalize( 

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

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

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

501 interval=aviz.MinMaxInterval(), 

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

503 ) 

504 else: 

505 norm = aviz.ImageNormalize( 

506 data, 

507 interval=aviz.MinMaxInterval(), 

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

509 ) 

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

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

512 x_line = 1 

513 y_line = 1 

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

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

516 ax.axis("off") 

517 if name is not None: 

518 ax.set_title(name) 

519 

520 try: 

521 len_sizes = len(sizes) 

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

523 if len_sizes == 1: 

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

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

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

527 else: 

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

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

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

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

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

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

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

535 if source is not None: 

536 _annotate_image(fig, source, len_sizes) 

537 

538 output = io.BytesIO() 

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

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

541 finally: 

542 plt.close(fig) 

543 

544 return output 

545 

546 

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

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

549_FLAG_GROUPS = { 

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

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

552 "forced": ["forced_PsfFlux_flag", "forced_PsfFlux_flag_noGoodPixels", 

553 "forced_PsfFlux_flag_edge"], 

554 "edge": ["pixelFlags_edge"], 

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

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

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

558 "bad": ["pixelFlags_bad"], 

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

560 "centroid": ["centroid_flag"], 

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

562 "shape_flag_parent_source"], 

563} 

564 

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

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

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

568# or a callable returning a bool. 

569_FLAG_TAGS = [ 

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

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

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

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

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

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

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

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

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

579] 

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

581 

582 

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

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

585 

586 Parameters 

587 ---------- 

588 fig : `matplotlib.Figure` 

589 Figure to be annotated. 

590 source : `numpy.record` 

591 DiaSource record of the object being plotted. 

592 len_sizes : `int` 

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

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

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

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

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

598 """ 

599 flag_color = "red" 

600 text_color = "grey" 

601 

602 if heights is None: 

603 if len_sizes == 1: 

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

605 else: 

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

607 

608 def label_color(group_key): 

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

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

611 

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

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

614 rows = [ 

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

616 [ 

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

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

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

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

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

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

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

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

625 ], 

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

627 [ 

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

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

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

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

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

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

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

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

636 ], 

637 # Row 2: PSF flux. 

638 [ 

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

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

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

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

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

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

645 ], 

646 # Row 3: aperture flux. 

647 [ 

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

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

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

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

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

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

654 ], 

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

656 [ 

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

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

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

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

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

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

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

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

665 ], 

666 ] 

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

668 for x, text, kwargs in row: 

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

670 

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

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

673 if callable(predicate): 

674 triggered = predicate(source) 

675 else: 

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

677 if triggered: 

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

679 

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

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

682 rb = source['reliability'] 

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

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

685 fontweight="bold") 

686 

687 

688class CutoutPath: 

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

690 

691 Supports local files, and id-chunked directories. 

692 

693 Parameters 

694 ---------- 

695 root : `str` 

696 Root file path to manage. 

697 chunk_size : `int`, optional 

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

699 subdirectory : `str`, optional 

700 Name of the subdirectory 

701 

702 Raises 

703 ------ 

704 RuntimeError 

705 Raised if chunk_size is not a power of 10. 

706 """ 

707 

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

709 self._root = root 

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

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

712 self._chunk_size = chunk_size 

713 self._subdirectory = subdirectory 

714 

715 def directory(self, id): 

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

717 

718 Parameters 

719 ---------- 

720 id : `int` 

721 Source id to create the path for. 

722 

723 Returns 

724 ------- 

725 directory: `str` 

726 Directory for this file. 

727 """ 

728 

729 def chunker(id, size): 

730 return (id // size)*size 

731 

732 if self._chunk_size is not None: 

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

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

735 else: 

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

737 

738 def __call__(self, id, filename): 

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

740 

741 Parameters 

742 ---------- 

743 id : `int` 

744 Source id to create the path for. 

745 filename: `str` 

746 Filename to write. 

747 

748 Returns 

749 ------- 

750 path : `str` 

751 Full path to the requested file. 

752 """ 

753 

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

755 

756 def exists(self, id, filename): 

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

758 

759 Parameters 

760 ---------- 

761 id : `int` 

762 Source id to create the path for. 

763 filename: `str` 

764 Filename to write. 

765 

766 Returns 

767 ------- 

768 exists : `bool` 

769 Does the supplied filename exist? 

770 """ 

771 

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

773 

774 def mkdir(self, id): 

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

776 

777 Parameters 

778 ---------- 

779 id : `int` 

780 Source id to create the path for. 

781 """ 

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

783 

784 

785def build_argparser(): 

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

787 script. 

788 

789 Returns 

790 ------- 

791 argparser : `argparse.ArgumentParser` 

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

793 command-line interface. 

794 """ 

795 parser = argparse.ArgumentParser( 

796 description=__doc__, 

797 formatter_class=argparse.RawDescriptionHelpFormatter, 

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

799 ) 

800 

801 apdbArgs = parser.add_mutually_exclusive_group(required=True) 

802 apdbArgs.add_argument( 

803 "--sqlitefile", 

804 default=None, 

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

806 ) 

807 apdbArgs.add_argument( 

808 "--namespace", 

809 default=None, 

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

811 " required for postgres connections." 

812 ) 

813 

814 parser.add_argument( 

815 "--postgres_url", 

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

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

818 ) 

819 

820 parser.add_argument( 

821 "--limit", 

822 default=5, 

823 type=int, 

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

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

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

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

828 ) 

829 parser.add_argument( 

830 "--all", 

831 default=False, 

832 action="store_true", 

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

834 ) 

835 

836 parser.add_argument( 

837 "-j", 

838 "--jobs", 

839 default=0, 

840 type=int, 

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

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

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

844 ) 

845 

846 parser.add_argument( 

847 "-C", 

848 "--configFile", 

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

850 ) 

851 parser.add_argument( 

852 "--collections", 

853 nargs="*", 

854 help=( 

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

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

857 "which may be very slow." 

858 ), 

859 ) 

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

861 parser.add_argument( 

862 "outputPath", 

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

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

865 ) 

866 parser.add_argument( 

867 "--reliabilityMin", 

868 type=float, 

869 default=None, 

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

871 ) 

872 parser.add_argument( 

873 "--reliabilityMax", 

874 type=float, 

875 default=None, 

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

877 ) 

878 return parser 

879 

880 

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

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

883 

884 Parameters 

885 ---------- 

886 sqlitefile : `str`, optional 

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

888 postgres_url : `str`, optional 

889 Postgres connection URL to connect to APDB. 

890 namespace : `str`, optional 

891 Postgres schema to load from; required with postgres_url. 

892 

893 Returns 

894 ------- 

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

896 Query instance to use to load data from APDB. 

897 

898 Raises 

899 ------ 

900 RuntimeError 

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

902 """ 

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

904 apdb_query = apdb.ApdbSqliteQuery(sqlitefile) 

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

906 apdb_query = apdb.ApdbPostgresQuery(namespace, postgres_url) 

907 else: 

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

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

910 return apdb_query 

911 

912 

913def run_cutouts(args): 

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

915 

916 Parameters 

917 ---------- 

918 args : `argparse.Namespace` 

919 Parsed commandline arguments. 

920 """ 

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

922 logging.basicConfig( 

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

924 ) 

925 

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

927 apdb_query = _make_apdbQuery(sqlitefile=args.sqlitefile, 

928 postgres_url=args.postgres_url, 

929 namespace=args.namespace) 

930 

931 config = PlotImageSubtractionCutoutsConfig() 

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

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

934 config.freeze() 

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

936 

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

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

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

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

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

942 "apFlux", "apFluxErr", "snr", "psfFlux", "psfFluxErr", 

943 "isDipole", "trailLength", "band", "extendedness", 

944 "pixelFlags_bad", "pixelFlags_cr", "pixelFlags_crCenter", 

945 "pixelFlags_edge", "pixelFlags_interpolated", "pixelFlags_interpolatedCenter", 

946 "pixelFlags_offimage", "pixelFlags_saturated", "pixelFlags_saturatedCenter", 

947 "pixelFlags_suspect", "pixelFlags_suspectCenter", "pixelFlags_streak", 

948 "pixelFlags_streakCenter", "pixelFlags_injected", "pixelFlags_injectedCenter", 

949 "pixelFlags_injected_template", "pixelFlags_injected_templateCenter"] 

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

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

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

953 

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

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

956 if not args.all: 

957 data = next(getter) 

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

959 else: 

960 sources = [] 

961 count = apdb_query.count_sources() 

962 for i, data in enumerate(getter): 

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

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

965 cutouts.write_manifest(sources) 

966 

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

968 # Write a dataframe with only diasources successfully written. 

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

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

971 

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

973 

974 

975def main(): 

976 args = build_argparser().parse_args() 

977 run_cutouts(args)