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
« 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/>.
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
44from lsst.images import Box, NoOverlapError
45import numpy as np
46import pandas as pd
48from . import apdb
49from .imageUtils import as_visit_image, cutout, cutout_box, get_visit_image, pixel_scale
51_log = logging.getLogger(__name__)
54class _ButlerCache:
55 """Global class to handle butler queries, to allow lru_cache and
56 `multiprocessing.Pool` to work together.
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 """
62 def set(self, butler, config):
63 """Call this to store a Butler and Config instance before using the
64 global class instance.
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()
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.
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.
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
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")
118 return science, template, difference
120 @functools.lru_cache(maxsize=4)
121 def get_catalog(self, instrument, detector, visit):
122 """Return the diaSrc catalog from the butler.
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.
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)
142# Global used within each multiprocessing worker (or single process).
143butler_cache = _ButlerCache()
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 )
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.
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"
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')
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
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.
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.
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
251 def write_manifest(self, sources):
252 """Save a Zooniverse manifest attaching image URLs to source ids.
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.")
265 def _make_manifest(self, sources):
266 """Return a Zooniverse manifest attaching image URLs to source ids.
268 Parameters
269 ----------
270 sources : `list` [`int`]
271 The diaSourceIds of the sources that had cutouts succesfully made.
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
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.
289 Creates ``images/`` and ``numpy/`` subdirectories if they
290 do not already exist; images are written there as PNG and npy files.
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.
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")
312 # Exclude index if they are replicated in columns.
313 indexNotInColumns = not any(index in data.columns for index in data.index.names)
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))
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]
333 def _do_one_source(self, source):
334 """Make cutouts for one diaSource.
336 Parameters
337 ----------
338 source : `numpy.record`, optional
339 DiaSource record for this cutout, to add metadata to the image.
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()
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
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.
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``.
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]
451 return self._plot_cutout(cutout_science,
452 cutout_template,
453 cutout_difference,
454 scale,
455 sizes,
456 source=source)
458 def _plot_cutout(self, science, template, difference, scale, sizes, source=None):
459 """Plot the cutouts for a source in one image.
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.
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
492 # TODO DM-32014: how do we color masked pixels (including edges)?
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)
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)
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)
544 return output
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}
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.
583def _annotate_image(fig, source, len_sizes, heights=None):
584 """Annotate the cutouts image with metadata and flags.
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"
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]
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
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)
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")
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")
688class CutoutPath:
689 """Manage paths to image cutouts with filenames based on diaSourceId.
691 Supports local files, and id-chunked directories.
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
702 Raises
703 ------
704 RuntimeError
705 Raised if chunk_size is not a power of 10.
706 """
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
715 def directory(self, id):
716 """Return the directory to store the output in.
718 Parameters
719 ----------
720 id : `int`
721 Source id to create the path for.
723 Returns
724 -------
725 directory: `str`
726 Directory for this file.
727 """
729 def chunker(id, size):
730 return (id // size)*size
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}")
738 def __call__(self, id, filename):
739 """Return the full path to a diaSource cutout.
741 Parameters
742 ----------
743 id : `int`
744 Source id to create the path for.
745 filename: `str`
746 Filename to write.
748 Returns
749 -------
750 path : `str`
751 Full path to the requested file.
752 """
754 return os.path.join(self.directory(id), filename)
756 def exists(self, id, filename):
757 """Return True if the file already exists.
759 Parameters
760 ----------
761 id : `int`
762 Source id to create the path for.
763 filename: `str`
764 Filename to write.
766 Returns
767 -------
768 exists : `bool`
769 Does the supplied filename exist?
770 """
772 return os.path.exists(os.path.join(self.directory(id), filename))
774 def mkdir(self, id):
775 """Make the directory tree to write this cutout id to.
777 Parameters
778 ----------
779 id : `int`
780 Source id to create the path for.
781 """
782 os.makedirs(self.directory(id), exist_ok=True)
785def build_argparser():
786 """Construct an argument parser for the ``plotImageSubtractionCutouts``
787 script.
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 )
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 )
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 )
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 )
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 )
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
881def _make_apdbQuery(sqlitefile=None, postgres_url=None, namespace=None):
882 """Return a query connection to the specified APDB.
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.
893 Returns
894 -------
895 apdb_query : `lsst.analysis.ap.ApdbQuery`
896 Query instance to use to load data from APDB.
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
913def run_cutouts(args):
914 """Run PlotImageSubtractionCutoutsTask on the parsed commandline arguments.
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 )
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)
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)
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)
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)
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)
972 print(f"Generated {len(sources)} diaSource cutouts to {args.outputPath}.")
975def main():
976 args = build_argparser().parse_args()
977 run_cutouts(args)