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