Coverage for python/lsst/pipe/tasks/extended_psf/extended_psf_cutout.py: 60%
161 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-23 01:58 -0700
1# This file is part of pipe_tasks.
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__all__ = ["ExtendedPsfCutoutConnections", "ExtendedPsfCutoutConfig", "ExtendedPsfCutoutTask"]
24import astropy.units as u
25import numpy as np
26from astropy.coordinates import SkyCoord
27from astropy.table import Table
29from lsst.afw.cameraGeom import FIELD_ANGLE, FOCAL_PLANE, PIXELS
30from lsst.afw.detection import footprintsToNumpy
31from lsst.afw.geom import makeModifiedWcs
32from lsst.afw.geom.transformFactory import makeTransform
33from lsst.afw.image import ExposureF, MaskedImageF
34from lsst.afw.math import BackgroundList, WarpingControl, warpImage
35from lsst.afw.table import SourceCatalog
36from lsst.geom import (
37 AffineTransform,
38 Box2I,
39 Extent2D,
40 Extent2I,
41 Point2D,
42 Point2I,
43 SpherePoint,
44 arcseconds,
45 floor,
46 radians,
47)
48from lsst.images import (
49 GeneralFrame,
50 Image,
51 Mask,
52 MaskPlane,
53 SkyProjection,
54 get_legacy_visit_image_mask_planes,
55)
56from lsst.meas.algorithms import (
57 LoadReferenceObjectsConfig,
58 ReferenceObjectLoader,
59 WarpedPsf,
60)
61from lsst.pex.config import ChoiceField, ConfigField, Field, ListField
62from lsst.pipe.base import PipelineTask, PipelineTaskConfig, PipelineTaskConnections, Struct
63from lsst.pipe.base.connectionTypes import Input, Output, PrerequisiteInput
64from lsst.utils.timer import timeMethod
66from .extended_psf_candidates import ExtendedPsfCandidate, ExtendedPsfCandidateInfo, ExtendedPsfCandidates
68NEIGHBOR_MASK_PLANE = "NEIGHBOR"
71class ExtendedPsfCutoutConnections(
72 PipelineTaskConnections,
73 dimensions=("instrument", "visit", "detector"),
74):
75 """Connections for ExtendedPsfCutoutTask."""
77 ref_cat = PrerequisiteInput(
78 name="the_monster_20250219",
79 storageClass="SimpleCatalog",
80 doc="Reference catalog that contains star positions.",
81 dimensions=("skypix",),
82 multiple=True,
83 deferLoad=True,
84 )
85 input_exposure = Input(
86 name="preliminary_visit_image",
87 storageClass="ExposureF",
88 doc="Background-subtracted input exposure from which to extract cutouts around a star.",
89 dimensions=("visit", "detector"),
90 )
91 input_background = Input(
92 name="preliminary_visit_image_background",
93 storageClass="Background",
94 doc="Background model for the input exposure, to be added back on during processing.",
95 dimensions=("visit", "detector"),
96 )
97 input_source_catalog = Input(
98 name="single_visit_star_footprints",
99 storageClass="SourceCatalog",
100 doc="Source catalog containing footprints on the input exposure, used to mask neighboring sources.",
101 dimensions=("visit", "detector"),
102 )
103 extended_psf_candidates = Output(
104 name="extended_psf_candidates",
105 storageClass="ExtendedPsfCandidates",
106 doc="Set of preprocessed cutouts, each centered on a single star.",
107 dimensions=("visit", "detector"),
108 )
111class ExtendedPsfCutoutConfig(
112 PipelineTaskConfig,
113 pipelineConnections=ExtendedPsfCutoutConnections,
114):
115 """Configuration parameters for ExtendedPsfCutoutTask."""
117 # Star selection
118 mag_range = ListField[float](
119 doc="Magnitude range in Gaia G. Cutouts will be made for all stars in this range.",
120 default=[10, 18],
121 )
122 exclude_arcsec_radius = Field[float](
123 doc="No cutouts will be generated for stars with a neighboring star in the range "
124 "``exclude_mag_range`` mag within ``exclude_arcsec_radius`` arcseconds.",
125 default=5,
126 )
127 exclude_mag_range = ListField[float](
128 doc="No cutouts will be generated for stars with a neighboring star in the range "
129 "``exclude_mag_range`` mag within ``exclude_arcsec_radius`` arcseconds.",
130 default=[0, 20],
131 )
132 min_area_fraction = Field[float](
133 doc="Minimum fraction of the cutout area, post-masking, that must remain for it to be retained.",
134 default=0.1,
135 )
136 bad_mask_planes = ListField[str](
137 doc="Mask planes that identify excluded pixels for the calculation of ``min_area_fraction``.",
138 default=[
139 "BAD",
140 "CR",
141 "CROSSTALK",
142 "EDGE",
143 "NO_DATA",
144 "SAT",
145 "SUSPECT",
146 "UNMASKEDNAN",
147 NEIGHBOR_MASK_PLANE,
148 ],
149 )
150 min_focal_plane_radius = Field[float](
151 doc="Minimum distance to the center of the focal plane, in mm. "
152 "Stars with a focal plane radius smaller than this will be omitted.",
153 default=0.0,
154 )
155 max_focal_plane_radius = Field[float](
156 doc="Maximum distance to the center of the focal plane, in mm. "
157 "Stars with a focal plane radius larger than this will be omitted.",
158 default=np.inf,
159 )
161 # Cutout geometry
162 cutout_size = ListField[int](
163 doc="Size of the cutouts to be extracted, in pixels.",
164 default=[251, 251],
165 )
166 warping_kernel_name = ChoiceField[str](
167 doc="Warping kernel for image data warping.",
168 default="lanczos5",
169 allowed={
170 "bilinear": "bilinear interpolation",
171 "lanczos3": "Lanczos kernel of order 3",
172 "lanczos4": "Lanczos kernel of order 4",
173 "lanczos5": "Lanczos kernel of order 5",
174 },
175 )
176 mask_warping_kernel_name = ChoiceField[str](
177 doc="Warping kernel for mask warping. Typically a more conservative kernel (e.g. with less ringing) "
178 "is desirable for warping masks than for warping image data.",
179 default="bilinear",
180 allowed={
181 "bilinear": "bilinear interpolation",
182 "lanczos3": "Lanczos kernel of order 3",
183 "lanczos4": "Lanczos kernel of order 4",
184 "lanczos5": "Lanczos kernel of order 5",
185 },
186 )
188 # Misc
189 load_reference_objects_config = ConfigField[LoadReferenceObjectsConfig](
190 doc="Reference object loader for astrometric calibration.",
191 )
192 ref_cat_filter_name = Field[str](
193 doc="Name of the filter in the reference catalog to use for star selection. ",
194 default="phot_g_mean",
195 )
198class ExtendedPsfCutoutTask(PipelineTask):
199 """Extract extended PSF cutouts, and warp to the same pixel grid.
201 The ExtendedPsfCutoutTask is used to extract, process, and store small
202 image cutouts around stars.
203 This task essentially consists of two principal steps.
204 First, it identifies stars within an exposure using a reference
205 catalog and extracts a cutout around each.
206 Second, it shifts and warps each cutout to remove optical distortions and
207 sample all stars on the same pixel grid.
208 """
210 ConfigClass = ExtendedPsfCutoutConfig
211 _DefaultName = "extendedPsfCutout"
212 config: ExtendedPsfCutoutConfig
214 def runQuantum(self, butlerQC, inputRefs, outputRefs):
215 inputs = butlerQC.get(inputRefs)
216 ref_obj_loader = ReferenceObjectLoader(
217 dataIds=[ref.datasetRef.dataId for ref in inputRefs.ref_cat],
218 refCats=inputs.pop("ref_cat"),
219 name=self.config.connections.ref_cat,
220 config=self.config.load_reference_objects_config,
221 )
222 output = self.run(**inputs, ref_obj_loader=ref_obj_loader)
223 # Only ingest if output exists; prevents ingesting an empty FITS file
224 if output:
225 butlerQC.put(output, outputRefs)
227 @timeMethod
228 def run(
229 self,
230 input_exposure: ExposureF,
231 input_background: BackgroundList,
232 input_source_catalog: SourceCatalog,
233 ref_obj_loader: ReferenceObjectLoader,
234 ):
235 """Identify stars within an exposure using a reference catalog,
236 extract cutouts around each and warp/shift cutouts onto a common frame.
238 Parameters
239 ----------
240 input_exposure : `~lsst.afw.image.ExposureF`
241 The background-subtracted image to extract cutouts around stars.
242 input_background : `~lsst.afw.math.BackgroundList`
243 The background model associated with the input exposure.
244 input_source_catalog : `~lsst.afw.table.SourceCatalog`
245 The source catalog containing footprints on the input exposure.
246 ref_obj_loader : `~lsst.meas.algorithms.ReferenceObjectLoader`
247 Loader to find objects within a reference catalog.
249 Returns
250 -------
251 extended_psf_candidates :
252 `~lsst.pipe.tasks.extendedPsf.ExtendedPsfCandidates`
253 A set of cutouts, each centered on an extended PSF candidate.
254 """
255 extended_psf_candidate_table = self._get_extended_psf_candidate_table(ref_obj_loader, input_exposure)
257 extended_psf_candidates = self._get_extended_psf_candidates(
258 input_exposure,
259 input_background,
260 input_source_catalog,
261 extended_psf_candidate_table,
262 )
264 return Struct(extended_psf_candidates=extended_psf_candidates)
266 def _get_extended_psf_candidate_table(
267 self,
268 ref_obj_loader: ReferenceObjectLoader,
269 input_exposure: ExposureF,
270 ) -> Table:
271 """Get a table of extended PSF candidates from the reference catalog.
273 Trim the reference catalog to only those objects within the exposure
274 bounding box.
275 Then, select stars based on the specified magnitude range,
276 isolation criteria, and optionally focal plane radius criteria.
277 Finally, add columns with pixel coordinates and focal plane coordinates
278 for each extended PSF candidate.
280 Parameters
281 ----------
282 ref_obj_loader : `~lsst.meas.algorithms.ReferenceObjectLoader`
283 Loader to find objects within a reference catalog.
284 input_exposure : `~lsst.afw.image.ExposureF`
285 The exposure for which extended PSF candidates are being selected.
287 Returns
288 -------
289 extended_psf_candidate_table : `~astropy.table.Table`
290 Table of extended PSF candidates within the exposure.
291 """
292 bbox = input_exposure.getBBox()
293 wcs = input_exposure.getWcs()
294 detector = input_exposure.detector
296 # Load all ref cat stars within the padded exposure bounding box
297 within_region = ref_obj_loader.loadPixelBox(bbox, wcs, self.config.ref_cat_filter_name)
298 ref_cat_full = within_region.refCat
299 flux_field: str = within_region.fluxField
300 exclude_arcsec_radius = self.config.exclude_arcsec_radius * u.arcsec
302 # Convert mag ranges to flux in nJy for comparison with ref cat fluxes
303 flux_range_candidate = sorted(((self.config.mag_range * u.ABmag).to(u.nJy)).to_value())
304 flux_range_neighbor = sorted(((self.config.exclude_mag_range * u.ABmag).to(u.nJy)).to_value())
306 # Create a subset of ref cat stars that includes all stars that could
307 # potentially be either a candidate or a neighbor based on flux
308 flux_min = np.min((flux_range_candidate[0], flux_range_neighbor[0]))
309 flux_max = np.max((flux_range_candidate[1], flux_range_neighbor[1]))
310 maximal_subset = (ref_cat_full[flux_field] >= flux_min) & (ref_cat_full[flux_field] <= flux_max)
311 ref_cat_subset_columns = ("id", "coord_ra", "coord_dec", flux_field)
312 ref_cat_subset = Table(ref_cat_full.extract(*ref_cat_subset_columns, where=maximal_subset))
313 flux_subset = ref_cat_subset[flux_field]
315 # Identify candidate stars and their neighbors based on flux
316 is_candidate = (flux_subset >= flux_range_candidate[0]) & (flux_subset <= flux_range_candidate[1])
317 is_neighbor = (flux_subset >= flux_range_neighbor[0]) & (flux_subset <= flux_range_neighbor[1])
319 # Trim star coordinates to candidate and neighbor subsets
320 coords = SkyCoord(ref_cat_subset["coord_ra"], ref_cat_subset["coord_dec"], unit="rad")
321 coords_candidate = coords[is_candidate]
322 coords_neighbor = coords[is_neighbor]
324 # Identify candidate stars that have no contaminant neighbors
325 is_candidate_isolated = np.ones(len(coords_candidate), dtype=bool)
326 if len(coords_neighbor) > 0:
327 _, indices_candidate, angular_separation, _ = coords_candidate.search_around_sky(
328 coords_neighbor, exclude_arcsec_radius
329 )
330 indices_candidate = indices_candidate[angular_separation > 0 * u.arcsec] # Exclude self-matches
331 is_candidate_isolated[indices_candidate] = False
333 # Trim ref cat subset to isolated stars; add ancillary data
334 extended_psf_candidate_table = ref_cat_subset[is_candidate][is_candidate_isolated]
336 flux_nanojansky = extended_psf_candidate_table[flux_field][:] * u.nJy
337 extended_psf_candidate_table["mag"] = flux_nanojansky.to(u.ABmag).to_value() # AB magnitudes
339 zip_ra_dec = zip(
340 extended_psf_candidate_table["coord_ra"] * radians,
341 extended_psf_candidate_table["coord_dec"] * radians,
342 )
343 sphere_points = [SpherePoint(ra, dec) for ra, dec in zip_ra_dec]
344 pixel_coords = wcs.skyToPixel(sphere_points)
345 extended_psf_candidate_table["pixel_x"] = [pixel_coord.x for pixel_coord in pixel_coords]
346 extended_psf_candidate_table["pixel_y"] = [pixel_coord.y for pixel_coord in pixel_coords]
348 mm_coords = detector.transform(pixel_coords, PIXELS, FOCAL_PLANE)
349 mm_coords_x = np.array([mm_coord.x for mm_coord in mm_coords])
350 mm_coords_y = np.array([mm_coord.y for mm_coord in mm_coords])
351 radius_mm = np.sqrt(mm_coords_x**2 + mm_coords_y**2)
352 angle_radians = np.arctan2(mm_coords_y, mm_coords_x)
353 extended_psf_candidate_table["radius_mm"] = radius_mm
354 extended_psf_candidate_table["angle_radians"] = angle_radians
356 # Trim star catalog to those within the exposure bounding box,
357 # and optionally within a range of focal plane radii
358 within_bbox = extended_psf_candidate_table["pixel_x"] >= bbox.getMinX()
359 within_bbox &= extended_psf_candidate_table["pixel_x"] <= bbox.getMaxX()
360 within_bbox &= extended_psf_candidate_table["pixel_y"] >= bbox.getMinY()
361 within_bbox &= extended_psf_candidate_table["pixel_y"] <= bbox.getMaxY()
362 within_radii = extended_psf_candidate_table["radius_mm"] >= self.config.min_focal_plane_radius
363 within_radii &= extended_psf_candidate_table["radius_mm"] <= self.config.max_focal_plane_radius
364 extended_psf_candidate_table = extended_psf_candidate_table[within_bbox & within_radii]
366 self.log.info(
367 "Identified %i reference star%s in the field of view after applying magnitude and isolation "
368 "cuts.",
369 len(extended_psf_candidate_table),
370 "s" if len(extended_psf_candidate_table) != 1 else "",
371 )
373 return extended_psf_candidate_table
375 def _get_extended_psf_candidates(
376 self,
377 input_exposure: ExposureF,
378 input_background: BackgroundList | None,
379 footprints: SourceCatalog | np.ndarray,
380 extended_psf_candidate_table: Table,
381 ) -> ExtendedPsfCandidates | None:
382 """Extract and warp extended PSF candidate cutouts.
384 For each extended PSF candidate, extract a cutout from the input
385 exposure centered on the candidate's pixel coordinates.
386 Then, shift and warp the cutout to recenter on the candidate and align
387 each to the same orientation.
388 Finally, check the fraction of the cutout area that is masked
389 (e.g. due to neighboring sources or bad pixels), and only retain those
390 with sufficient unmasked area.
392 Parameters
393 ----------
394 input_exposure : `~lsst.afw.image.ExposureF`
395 The science image to extract extended PSF cutouts.
396 input_background : `~lsst.afw.math.BackgroundList` | None
397 The background model associated with the input exposure.
398 If provided, this will be added back on to the input image.
399 footprints : `~lsst.afw.table.SourceCatalog` | `numpy.ndarray`
400 The source catalog containing footprints on the input exposure, or
401 a 2D numpy array with the same dimensions as the input exposure
402 where each pixel value corresponds to the source footprint ID.
403 extended_psf_candidate_table : `~astropy.table.Table`
404 Table of extended PSF candidates for which to extract cutouts.
406 Returns
407 -------
408 extended_psf_candidates :
409 `~lsst.pipe.tasks.extendedPsf.ExtendedPsfCandidates` | None
410 A set of cutouts, each centered on an extended PSF candidate.
411 If no cutouts are retained post-masking, returns `None`.
412 """
413 warp_control = WarpingControl(self.config.warping_kernel_name, self.config.mask_warping_kernel_name)
414 bbox = input_exposure.getBBox()
416 # Prepare data: add bg back on, and convert to nJy
417 input_MI = input_exposure.getMaskedImage()
418 if input_background is not None: 418 ↛ 419line 418 didn't jump to line 419 because the condition on line 418 was never true
419 input_MI += input_background.getImage()
420 input_MI = input_exposure.photoCalib.calibrateImage(input_MI) # to nJy
422 # Generate unique footprint IDs for NEIGHBOR masking
423 input_MI.mask.addMaskPlane(NEIGHBOR_MASK_PLANE)
424 if isinstance(footprints, SourceCatalog): 424 ↛ 425line 424 didn't jump to line 425 because the condition on line 424 was never true
425 footprints = footprintsToNumpy(footprints, bbox, asBool=False)
427 # Establish pixel-to-boresight-pseudopixel transform
428 pixel_scale = input_exposure.wcs.getPixelScale(bbox.getCenter()).asArcseconds() * arcseconds
429 pixels_to_boresight_pseudopixels = input_exposure.detector.getTransform(PIXELS, FIELD_ANGLE).then(
430 makeTransform(AffineTransform.makeScaling(1 / pixel_scale.asRadians()))
431 )
433 # Cutout bounding boxes
434 cutout_radius = floor(Extent2D(*self.config.cutout_size) / 2)
435 cutout_bbox = Box2I(Point2I(0, 0), Extent2I(1, 1)).dilatedBy(
436 cutout_radius
437 ) # always odd, centered 0,0
438 cutout_radius_padded = floor((Extent2D(*self.config.cutout_size) * 1.42) / 2) # max possible req. pad
439 cutout_bbox_padded = Box2I(Point2I(0, 0), Extent2I(1, 1)).dilatedBy(cutout_radius_padded)
441 cutouts = []
442 focal_plane_radii_mm = []
443 for candidate in extended_psf_candidate_table:
444 pix_coord = Point2D(candidate["pixel_x"], candidate["pixel_y"])
446 # Set NEIGHBOR mask plane for all sources except the current one
447 neighbor_bit_mask = input_MI.mask.getPlaneBitMask(NEIGHBOR_MASK_PLANE)
448 input_MI.mask.clearMaskPlane(input_MI.mask.getMaskPlaneDict()[NEIGHBOR_MASK_PLANE])
449 candidate_id = footprints[int(pix_coord.y), int(pix_coord.x)]
450 neighbor_mask = (footprints != 0) & (footprints != candidate_id)
451 input_MI.mask.array[neighbor_mask] |= neighbor_bit_mask
453 # Define linear shifting and rotation to recenter and align cutouts
454 boresight_pseudopixel_coord = pixels_to_boresight_pseudopixels.applyForward(pix_coord)
455 shift = makeTransform(AffineTransform(Point2D(0, 0) - boresight_pseudopixel_coord))
456 rotation = makeTransform(AffineTransform.makeRotation(-candidate["angle_radians"] * radians))
457 pixels_to_cutout_frame = pixels_to_boresight_pseudopixels.then(shift).then(rotation)
459 # Warp the image and mask to the cutout frame
460 cutout_MI = MaskedImageF(cutout_bbox_padded)
461 warpImage(cutout_MI, input_MI, pixels_to_cutout_frame, warp_control)
462 cutout_MI = cutout_MI[cutout_bbox]
464 # Skip if masked area fraction is too high
465 bad_bit_mask = cutout_MI.mask.getPlaneBitMask(self.config.bad_mask_planes)
466 good = (cutout_MI.mask.array & bad_bit_mask) == 0
467 good_frac = np.sum(good) / cutout_MI.mask.array.size
468 if good_frac < self.config.min_area_fraction: 468 ↛ 469line 468 didn't jump to line 469 because the condition on line 468 was never true
469 continue
471 # Define a WCS for the cutout consistent with the warping
472 cutout_wcs = makeModifiedWcs(pixels_to_cutout_frame, input_exposure.wcs, False)
473 sky_projection = SkyProjection.from_legacy(cutout_wcs, GeneralFrame(unit=u.pixel))
475 # Compute the kernel image of the PSF at the cutout center
476 psf_warped = WarpedPsf(input_exposure.getPsf(), pixels_to_cutout_frame, warp_control)
477 psf_kernel_image = Image.from_legacy(psf_warped.computeKernelImage(Point2D(0, 0)))
479 # Assemble the star info to be persisted alongside the image data
480 star_info = ExtendedPsfCandidateInfo(
481 visit=input_exposure.visitInfo.getId(),
482 detector=input_exposure.detector.getId(),
483 ref_id=candidate["id"],
484 ref_mag=candidate["mag"],
485 position_x=pix_coord.x,
486 position_y=pix_coord.y,
487 focal_plane_radius=candidate["radius_mm"] * u.mm,
488 focal_plane_angle=candidate["angle_radians"] * u.rad,
489 )
491 plane_map = get_legacy_visit_image_mask_planes()
492 plane_map["NEIGHBOR"] = MaskPlane(
493 "NEIGHBOR", "Flux in pixel is attributed to a neighboring source detection footprint."
494 )
496 # Generate an extended PSF candidate and store outputs
497 cutout = ExtendedPsfCandidate(
498 image=Image.from_legacy(cutout_MI.image),
499 mask=Mask.from_legacy(cutout_MI.mask, plane_map=plane_map),
500 variance=Image.from_legacy(cutout_MI.variance),
501 sky_projection=sky_projection,
502 psf_kernel_image=psf_kernel_image,
503 star_info=star_info,
504 )
505 cutouts.append(cutout)
506 focal_plane_radii_mm.append(candidate["radius_mm"])
508 num_stars = len(extended_psf_candidate_table)
509 num_excluded = num_stars - len(cutouts)
510 percent_excluded = 100.0 * num_excluded / num_stars if num_stars > 0 else 0.0
511 self.log.info(
512 "Extracted %i extended PSF candidate%s. "
513 "Excluded %i star%s (%.1f%%) with an unmasked area fraction below %s.",
514 len(cutouts),
515 "" if len(cutouts) == 1 else "s",
516 num_excluded,
517 "" if num_excluded == 1 else "s",
518 percent_excluded,
519 self.config.min_area_fraction,
520 )
522 if not cutouts: 522 ↛ 523line 522 didn't jump to line 523 because the condition on line 522 was never true
523 self.log.warning(
524 "No extended PSF candidates were retained from %i selected reference star%s.",
525 num_stars,
526 "" if num_stars == 1 else "s",
527 )
528 return None
530 metadata = {
531 "FOCAL_PLANE_RADIUS_MM_MIN": np.min(focal_plane_radii_mm),
532 "FOCAL_PLANE_RADIUS_MM_MAX": np.max(focal_plane_radii_mm),
533 }
534 return ExtendedPsfCandidates(cutouts, metadata=metadata)