Coverage for python/lsst/analysis/ap/plotImageSubtractionCutouts.py: 87%
289 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-09-29 10:31 +0000
« prev ^ index » next coverage.py v7.16.2, created at 2026-09-29 10:31 +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/>.
22"""Construct template/image/difference cutouts for upload to Zooniverse, or
23to just to view as images.
24"""
26__all__ = ["PlotImageSubtractionCutoutsConfig", "PlotImageSubtractionCutoutsTask", "CutoutPath"]
28import argparse
29import functools
30import io
31import logging
32import multiprocessing
33import os
34from math import log10
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
47from . import apdb
49_log = logging.getLogger(__name__)
52class _ButlerCache:
53 """Global class to handle butler queries, to allow lru_cache and
54 `multiprocessing.Pool` to work together.
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 """
60 def set(self, butler, config):
61 """Call this to store a Butler and Config instance before using the
62 global class instance.
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()
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.
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.
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
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)
112 return science, template, difference
114 @functools.lru_cache(maxsize=4)
115 def get_catalog(self, instrument, detector, visit):
116 """Return the diaSrc catalog from the butler.
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.
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)
136# Global used within each multiprocessing worker (or single process).
137butler_cache = _ButlerCache()
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 )
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.
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"
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')
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
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.
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.
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
245 def write_manifest(self, sources):
246 """Save a Zooniverse manifest attaching image URLs to source ids.
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.")
259 def _make_manifest(self, sources):
260 """Return a Zooniverse manifest attaching image URLs to source ids.
262 Parameters
263 ----------
264 sources : `list` [`int`]
265 The diaSourceIds of the sources that had cutouts succesfully made.
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
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.
283 Creates ``images/`` and ``numpy/`` subdirectories if they
284 do not already exist; images are written there as PNG and npy files.
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.
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")
306 # Exclude index if they are replicated in columns.
307 indexNotInColumns = not any(index in data.columns for index in data.index.names)
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))
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]
327 def _do_one_source(self, source):
328 """Make cutouts for one diaSource.
330 Parameters
331 ----------
332 source : `numpy.record`, optional
333 DiaSource record for this cutout, to add metadata to the image.
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()
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
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.
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``.
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]
442 return self._plot_cutout(cutout_science,
443 cutout_template,
444 cutout_difference,
445 scale,
446 sizes,
447 source=source)
449 def _plot_cutout(self, science, template, difference, scale, sizes, source=None):
450 """Plot the cutouts for a source in one image.
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.
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
483 # TODO DM-32014: how do we color masked pixels (including edges)?
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)
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)
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)
535 return output
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}
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.
574def _annotate_image(fig, source, len_sizes, heights=None):
575 """Annotate the cutouts image with metadata and flags.
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"
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]
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
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)
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")
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")
679class CutoutPath:
680 """Manage paths to image cutouts with filenames based on diaSourceId.
682 Supports local files, and id-chunked directories.
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
693 Raises
694 ------
695 RuntimeError
696 Raised if chunk_size is not a power of 10.
697 """
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
706 def directory(self, id):
707 """Return the directory to store the output in.
709 Parameters
710 ----------
711 id : `int`
712 Source id to create the path for.
714 Returns
715 -------
716 directory: `str`
717 Directory for this file.
718 """
720 def chunker(id, size):
721 return (id // size)*size
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}")
729 def __call__(self, id, filename):
730 """Return the full path to a diaSource cutout.
732 Parameters
733 ----------
734 id : `int`
735 Source id to create the path for.
736 filename: `str`
737 Filename to write.
739 Returns
740 -------
741 path : `str`
742 Full path to the requested file.
743 """
745 return os.path.join(self.directory(id), filename)
747 def exists(self, id, filename):
748 """Return True if the file already exists.
750 Parameters
751 ----------
752 id : `int`
753 Source id to create the path for.
754 filename: `str`
755 Filename to write.
757 Returns
758 -------
759 exists : `bool`
760 Does the supplied filename exist?
761 """
763 return os.path.exists(os.path.join(self.directory(id), filename))
765 def mkdir(self, id):
766 """Make the directory tree to write this cutout id to.
768 Parameters
769 ----------
770 id : `int`
771 Source id to create the path for.
772 """
773 os.makedirs(self.directory(id), exist_ok=True)
776def build_argparser():
777 """Construct an argument parser for the ``plotImageSubtractionCutouts``
778 script.
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 )
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 )
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 )
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 )
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 )
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
872def _make_apdbQuery(sqlitefile=None, postgres_url=None, namespace=None):
873 """Return a query connection to the specified APDB.
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.
884 Returns
885 -------
886 apdb_query : `lsst.analysis.ap.ApdbQuery`
887 Query instance to use to load data from APDB.
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
904def run_cutouts(args):
905 """Run PlotImageSubtractionCutoutsTask on the parsed commandline arguments.
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 )
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)
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)
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)
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)
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)
963 print(f"Generated {len(sources)} diaSource cutouts to {args.outputPath}.")
966def main():
967 args = build_argparser().parse_args()
968 run_cutouts(args)