Coverage for python/lsst/pipe/tasks/extended_psf/extended_psf_subtract.py: 64%
195 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-25 09:10 +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__ = [
23 "ExtendedPsfSubtractConnections",
24 "ExtendedPsfSubtractConfig",
25 "ExtendedPsfSubtractTask",
26]
28import astropy.units as u
29import numpy as np
30from astropy.coordinates import SkyCoord
31from astropy.table import Table
33from lsst.afw.cameraGeom import FIELD_ANGLE, FOCAL_PLANE, PIXELS
34from lsst.afw.geom.transformFactory import makeTransform
35from lsst.afw.image import ExposureF, ImageF
36from lsst.afw.math import BackgroundList, WarpingControl, warpImage
37from lsst.geom import (
38 AffineTransform,
39 Box2I,
40 Extent2I,
41 Point2D,
42 Point2I,
43 SpherePoint,
44 arcseconds,
45 radians,
46)
47from lsst.meas.algorithms import (
48 LoadReferenceObjectsConfig,
49 ReferenceObjectLoader,
50 SubtractBackgroundTask,
51)
52from lsst.pex.config import (
53 ChoiceField,
54 ConfigField,
55 ConfigurableField,
56 Field,
57 ListField,
58)
59from lsst.pipe.base import PipelineTask, PipelineTaskConfig, PipelineTaskConnections, Struct
60from lsst.pipe.base.connectionTypes import Input, Output, PrerequisiteInput
61from lsst.utils.timer import timeMethod
63from .extended_psf_image import ExtendedPsfImage
66class ExtendedPsfSubtractConnections(
67 PipelineTaskConnections,
68 dimensions=("instrument", "visit", "detector"),
69):
70 """Connections for ExtendedPsfSubtractTask."""
72 ref_cat = PrerequisiteInput(
73 name="the_monster_20250219",
74 storageClass="SimpleCatalog",
75 doc="Reference catalog that contains star positions.",
76 dimensions=("skypix",),
77 multiple=True,
78 deferLoad=True,
79 )
80 preliminary_visit_image = Input(
81 name="preliminary_visit_image",
82 storageClass="ExposureF",
83 doc="Input background-subtracted image.",
84 dimensions=("visit", "detector"),
85 )
86 preliminary_visit_image_background = Input(
87 name="preliminary_visit_image_background",
88 storageClass="Background",
89 doc=(
90 "Supplied input background model associated with preliminary_visit_image. "
91 "When do_restore_background is True, it is added back before fitting/subtracting the "
92 "extended PSF."
93 ),
94 dimensions=("visit", "detector"),
95 )
96 extended_psf = Input(
97 name="extended_psf",
98 storageClass="ExtendedPsfImage",
99 doc="Extended PSF model image.",
100 dimensions=("visit", "detector"),
101 )
102 preliminary_visit_image_extended_psf_subtracted = Output(
103 name="preliminary_visit_image_extended_psf_subtracted",
104 storageClass="ExposureF",
105 doc=(
106 "Preliminary-visit image with fitted extended-PSF subtracted. "
107 "Background may optionally be subtracted, depending on configuration."
108 ),
109 dimensions=("visit", "detector"),
110 )
111 preliminary_visit_image_extended_psf_subtracted_background = Output(
112 name="preliminary_visit_image_extended_psf_subtracted_background",
113 storageClass="Background",
114 doc="Background model re-estimated after extended-PSF subtraction.",
115 dimensions=("visit", "detector"),
116 )
118 def __init__(self, *, config=None):
119 super().__init__(config=config)
120 if config is not None and not config.do_restore_background:
121 del self.preliminary_visit_image_background
122 if config is not None and not config.do_rerun_background_subtraction:
123 del self.preliminary_visit_image_extended_psf_subtracted_background
126class ExtendedPsfSubtractConfig(
127 PipelineTaskConfig,
128 pipelineConnections=ExtendedPsfSubtractConnections,
129):
130 """Configuration parameters for ExtendedPsfSubtractTask."""
132 # Star selection
133 mag_range = ListField[float](
134 doc="Magnitude range in Gaia G for subtraction stars.",
135 default=[10, 18],
136 )
137 exclude_arcsec_radius = Field[float](
138 doc=(
139 "No subtraction fit will be done for stars that have a neighboring star in the "
140 "exclude_mag_range within exclude_arcsec_radius arcseconds."
141 ),
142 default=5,
143 )
144 exclude_mag_range = ListField[float](
145 doc="Magnitude range used when searching for neighboring contaminants.",
146 default=[0, 20],
147 )
148 min_focal_plane_radius = Field[float](
149 doc="Minimum focal-plane radius in mm for subtraction stars.",
150 default=0.0,
151 )
152 max_focal_plane_radius = Field[float](
153 doc="Maximum focal-plane radius in mm for subtraction stars.",
154 default=np.inf,
155 )
156 max_stars_per_detector = Field[int](
157 doc=(
158 "Maximum number of stars to subtract per detector; 0 means no cap. "
159 "If this limit is reached, the selected stars are truncated to the brightest "
160 "objects first (lowest magnitude values)."
161 ),
162 default=0,
163 )
165 # Warping
166 warping_kernel_name = ChoiceField[str](
167 doc="Warping kernel for model image 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 )
177 # Fitting
178 bad_mask_planes = ListField[str](
179 doc="Mask planes excluded during amplitude fitting.",
180 default=[
181 "BAD",
182 "COSMIC_RAY",
183 "CROSSTALK",
184 "DETECTION_EDGE",
185 "NO_DATA",
186 "SATURATED",
187 "SUSPECT",
188 "UNMASKED_NAN",
189 ],
190 )
191 min_model_value = Field[float](
192 doc=(
193 "Minimum warped-model pixel value allowed in the per-star amplitude fit. "
194 "Only pixels with model > min_model_value are used when solving for PSF scale, "
195 "after mask and variance filtering. Increasing this emphasizes the bright core and "
196 "suppresses noisy wings; decreasing it includes more wing pixels in the fit."
197 ),
198 default=0.0,
199 )
200 min_fit_pixels = Field[int](
201 doc="Minimum number of usable pixels required to fit a star.",
202 default=20,
203 )
204 bg_order = Field[int](
205 doc=(
206 "Order of polynomial to fit for local background around each star. "
207 "Set to 0 for a constant pedestal, 1 for a planar background, and higher values "
208 "for additional flexibility."
209 ),
210 default=1,
211 )
213 # Background
214 do_restore_background = Field[bool](
215 doc=(
216 "Restore the supplied preliminary_visit_image_background onto the working image before "
217 "star fitting/subtraction."
218 ),
219 default=True,
220 )
221 do_rerun_background_subtraction = Field[bool](
222 doc="Re-estimate and subtract background after extended-PSF subtraction.",
223 default=True,
224 )
225 subtract_background = ConfigurableField(
226 target=SubtractBackgroundTask,
227 doc="Subtask used to re-estimate background after extended-PSF subtraction.",
228 )
230 # Misc
231 load_reference_objects_config = ConfigField[LoadReferenceObjectsConfig](
232 doc="Reference object loader for astrometric lookup.",
233 )
234 ref_cat_filter_name = Field[str](
235 doc="Name of the filter in the reference catalog to use for star selection.",
236 default="phot_g_mean",
237 )
240class ExtendedPsfSubtractTask(PipelineTask):
241 """Subtract a fitted extended PSF model from stars in a detector image."""
243 ConfigClass = ExtendedPsfSubtractConfig
244 _DefaultName = "extendedPsfSubtract"
245 config: ExtendedPsfSubtractConfig
247 def __init__(self, initInputs=None, *args, **kwargs):
248 super().__init__(initInputs=initInputs, *args, **kwargs)
249 self._bg_powers = [
250 (i, j)
251 for i in range(self.config.bg_order + 1)
252 for j in range(self.config.bg_order + 1)
253 if i + j <= self.config.bg_order
254 ]
255 self.makeSubtask("subtract_background")
257 def runQuantum(self, butlerQC, inputRefs, outputRefs):
258 inputs = butlerQC.get(inputRefs)
259 ref_obj_loader = ReferenceObjectLoader(
260 dataIds=[ref.datasetRef.dataId for ref in inputRefs.ref_cat],
261 refCats=inputs.pop("ref_cat"),
262 name=self.config.connections.ref_cat,
263 config=self.config.load_reference_objects_config,
264 )
265 output = self.run(ref_obj_loader=ref_obj_loader, **inputs)
266 if output:
267 butlerQC.put(output, outputRefs)
269 @timeMethod
270 def run(
271 self,
272 preliminary_visit_image: ExposureF,
273 extended_psf: ExtendedPsfImage,
274 ref_obj_loader: ReferenceObjectLoader,
275 preliminary_visit_image_background: BackgroundList | None = None,
276 ):
277 """Subtract fitted extended-PSF models from selected stars.
279 This method clones the input exposure, optionally restores the
280 associated background model for fitting, selects subtraction stars from
281 the reference catalog, and fits/subtracts the warped extended-PSF
282 model for each star in sequence.
283 Stars are always processed in magnitude order from brightest to
284 faintest, and ``max_stars_per_detector`` (if non-zero) is applied after
285 sorting so only the brightest stars are retained.
286 After per-star subtraction, the method optionally re-estimates the
287 background model. When re-estimation is enabled, both updated
288 exposure and background outputs are returned.
290 Parameters
291 ----------
292 preliminary_visit_image : `lsst.afw.image.ExposureF`
293 Background-subtracted image.
294 extended_psf : `ExtendedPsfImage`
295 Extended PSF model to be warped and fit per star.
296 ref_obj_loader : `lsst.meas.algorithms.ReferenceObjectLoader`
297 Reference object loader used for star selection.
298 preliminary_visit_image_background :
299 `lsst.afw.math.BackgroundList`, optional
300 Supplied input background model associated with the input image.
301 This is restored onto the working exposure when
302 ``do_restore_background`` is True.
303 """
304 output_exposure = preliminary_visit_image.clone()
305 if self.config.do_restore_background: 305 ↛ 310line 305 didn't jump to line 310 because the condition on line 305 was always true
306 if preliminary_visit_image_background is None: 306 ↛ 307line 306 didn't jump to line 307 because the condition on line 306 was never true
307 raise RuntimeError("do_restore_background=True requires preliminary_visit_image_background.")
308 output_exposure.maskedImage += preliminary_visit_image_background.getImage()
310 star_table = self._get_subtraction_star_table(ref_obj_loader, output_exposure)
311 order = np.argsort(star_table["mag"])
312 star_table = star_table[order]
313 if self.config.max_stars_per_detector > 0 and len(star_table) > self.config.max_stars_per_detector: 313 ↛ 314line 313 didn't jump to line 314 because the condition on line 313 was never true
314 star_table = star_table[: self.config.max_stars_per_detector]
316 model_image_legacy = extended_psf.image.to_legacy(copy=False)
317 warp_control = WarpingControl(self.config.warping_kernel_name)
319 n_attempted = 0
320 n_subtracted = 0
321 amplitude_values = []
323 for star in star_table:
324 n_attempted += 1
325 amplitude = self._subtract_one_star(
326 output_exposure,
327 model_image_legacy,
328 warp_control,
329 star,
330 )
331 if np.isfinite(amplitude): 331 ↛ 323line 331 didn't jump to line 323 because the condition on line 331 was always true
332 n_subtracted += 1
333 amplitude_values.append(amplitude)
335 if self.config.do_rerun_background_subtraction: 335 ↛ 345line 335 didn't jump to line 345 because the condition on line 335 was always true
336 # Subtract background in place and retrieve the background model.
337 bg_result = self.subtract_background.run(
338 exposure=output_exposure,
339 background=None,
340 backgroundToPhotometricRatio=None,
341 stats=True,
342 )
343 output_background = bg_result.background
345 metadata = output_exposure.getMetadata()
346 metadata.set("EPSFSUB_ATTEMPTED", int(n_attempted))
347 metadata.set("EPSFSUB_SUBTRACTED", int(n_subtracted))
348 if amplitude_values: 348 ↛ 353line 348 didn't jump to line 353 because the condition on line 348 was always true
349 metadata.set("EPSFSUB_AMP_MIN", float(np.min(amplitude_values)))
350 metadata.set("EPSFSUB_AMP_MAX", float(np.max(amplitude_values)))
351 metadata.set("EPSFSUB_AMP_MED", float(np.median(amplitude_values)))
353 self.log.info(
354 "Extended-PSF subtraction attempted on %d stars; successfully subtracted %d.",
355 n_attempted,
356 n_subtracted,
357 )
359 if self.config.do_rerun_background_subtraction: 359 ↛ 364line 359 didn't jump to line 364 because the condition on line 359 was always true
360 return Struct(
361 preliminary_visit_image_extended_psf_subtracted=output_exposure,
362 preliminary_visit_image_extended_psf_subtracted_background=output_background,
363 )
364 return Struct(preliminary_visit_image_extended_psf_subtracted=output_exposure)
366 def _get_subtraction_star_table(
367 self,
368 ref_obj_loader: ReferenceObjectLoader,
369 exposure: ExposureF,
370 ) -> Table:
371 """Build a table of subtraction stars from the reference catalog.
373 Parameters
374 ----------
375 ref_obj_loader : `~lsst.meas.algorithms.ReferenceObjectLoader`
376 Loader used to query reference objects in the detector footprint.
377 exposure : `~lsst.afw.image.ExposureF`
378 Exposure used to define the detector bounding box, WCS, and focal
379 plane geometry for star selection.
381 Returns
382 -------
383 star_table : `~astropy.table.Table`
384 Table of selected subtraction stars after magnitude, isolation,
385 detector-footprint, and focal-plane-radius filtering.
386 Includes per-star pixel and focal-plane coordinates.
387 """
388 bbox = exposure.getBBox()
389 wcs = exposure.getWcs()
390 detector = exposure.detector
392 within_region = ref_obj_loader.loadPixelBox(bbox, wcs, self.config.ref_cat_filter_name)
393 ref_cat_full = within_region.refCat
394 flux_field: str = within_region.fluxField
395 exclude_arcsec_radius = self.config.exclude_arcsec_radius * u.arcsec
397 flux_range_candidate = sorted(((self.config.mag_range * u.ABmag).to(u.nJy)).to_value())
398 flux_range_neighbor = sorted(((self.config.exclude_mag_range * u.ABmag).to(u.nJy)).to_value())
400 flux_min = np.min((flux_range_candidate[0], flux_range_neighbor[0]))
401 flux_max = np.max((flux_range_candidate[1], flux_range_neighbor[1]))
402 maximal_subset = (ref_cat_full[flux_field] >= flux_min) & (ref_cat_full[flux_field] <= flux_max)
403 ref_cat_subset = Table(
404 ref_cat_full.extract(
405 "id",
406 "coord_ra",
407 "coord_dec",
408 flux_field,
409 where=maximal_subset,
410 )
411 )
412 flux_subset = ref_cat_subset[flux_field]
414 is_candidate = (flux_subset >= flux_range_candidate[0]) & (flux_subset <= flux_range_candidate[1])
415 is_neighbor = (flux_subset >= flux_range_neighbor[0]) & (flux_subset <= flux_range_neighbor[1])
417 coords = SkyCoord(ref_cat_subset["coord_ra"], ref_cat_subset["coord_dec"], unit="rad")
418 coords_candidate = coords[is_candidate]
419 coords_neighbor = coords[is_neighbor]
421 is_candidate_isolated = np.ones(len(coords_candidate), dtype=bool)
422 if len(coords_neighbor) > 0:
423 _, indices_candidate, angular_separation, _ = coords_candidate.search_around_sky(
424 coords_neighbor, exclude_arcsec_radius
425 )
426 indices_candidate = indices_candidate[angular_separation > 0 * u.arcsec]
427 is_candidate_isolated[indices_candidate] = False
429 star_table = ref_cat_subset[is_candidate][is_candidate_isolated]
430 flux_njy = star_table[flux_field][:] * u.nJy
431 star_table["mag"] = flux_njy.to(u.ABmag).to_value()
433 sphere_points = [
434 SpherePoint(ra * radians, dec * radians)
435 for ra, dec in zip(star_table["coord_ra"], star_table["coord_dec"])
436 ]
437 pixel_coords = wcs.skyToPixel(sphere_points)
438 star_table["pixel_x"] = [pixel_coord.x for pixel_coord in pixel_coords]
439 star_table["pixel_y"] = [pixel_coord.y for pixel_coord in pixel_coords]
441 mm_coords = detector.transform(pixel_coords, PIXELS, FOCAL_PLANE)
442 mm_x = np.array([mm_coord.x for mm_coord in mm_coords])
443 mm_y = np.array([mm_coord.y for mm_coord in mm_coords])
444 radius_mm = np.sqrt(mm_x**2 + mm_y**2)
445 star_table["radius_mm"] = radius_mm
446 star_table["angle_radians"] = np.arctan2(mm_y, mm_x)
448 within_bbox = star_table["pixel_x"] >= bbox.getMinX()
449 within_bbox &= star_table["pixel_x"] <= bbox.getMaxX()
450 within_bbox &= star_table["pixel_y"] >= bbox.getMinY()
451 within_bbox &= star_table["pixel_y"] <= bbox.getMaxY()
452 within_radii = star_table["radius_mm"] >= self.config.min_focal_plane_radius
453 within_radii &= star_table["radius_mm"] <= self.config.max_focal_plane_radius
454 star_table = star_table[within_bbox & within_radii]
456 return star_table
458 def _subtract_one_star(
459 self,
460 exposure: ExposureF,
461 model_image_legacy: ImageF,
462 warp_control: WarpingControl,
463 star: Table,
464 ) -> float:
465 """Warp, fit, and subtract one star model in detector coordinates.
467 Parameters
468 ----------
469 exposure : `~lsst.afw.image.ExposureF`
470 Exposure to update in place by subtracting the fitted PSF model.
471 model_image_legacy : `~lsst.afw.image.ImageF`
472 Empirical extended-PSF model image in the model frame.
473 warp_control : `~lsst.afw.math.WarpingControl`
474 Warping configuration used to map the model into detector
475 coordinates.
476 star : `~astropy.table.Row`
477 Row containing per-star subtraction metadata, including pixel
478 coordinates and focal-plane angle.
480 Returns
481 -------
482 amplitude : `float`
483 Fitted PSF amplitude if subtraction succeeds. Returns `numpy.nan`
484 when no valid fit is obtained or subtraction is rejected.
486 Notes
487 -----
488 This method modifies ``exposure`` in place for the selected local
489 bounding box by subtracting the fitted PSF model from the image plane.
490 """
491 detector_bbox = exposure.getBBox()
492 pix_coord = Point2D(star["pixel_x"], star["pixel_y"])
494 pixel_scale = exposure.getWcs().getPixelScale(detector_bbox.getCenter()).asArcseconds() * arcseconds
495 pixels_to_boresight_pseudopixels = exposure.detector.getTransform(PIXELS, FIELD_ANGLE).then(
496 makeTransform(AffineTransform.makeScaling(1 / pixel_scale.asRadians()))
497 )
499 boresight_pseudopixel_coord = pixels_to_boresight_pseudopixels.applyForward(pix_coord)
500 shift = makeTransform(AffineTransform(Point2D(0, 0) - boresight_pseudopixel_coord))
501 rotation = makeTransform(AffineTransform.makeRotation(-star["angle_radians"] * radians))
502 pixels_to_cutout_frame = pixels_to_boresight_pseudopixels.then(shift).then(rotation)
503 cutout_to_pixels = pixels_to_cutout_frame.inverted()
505 model_bbox = model_image_legacy.getBBox()
506 radius_x = model_bbox.getWidth() // 2
507 radius_y = model_bbox.getHeight() // 2
508 x_center = int(np.round(pix_coord.x))
509 y_center = int(np.round(pix_coord.y))
510 local_bbox = Box2I(
511 Point2I(x_center - radius_x, y_center - radius_y),
512 Extent2I(model_bbox.getWidth(), model_bbox.getHeight()),
513 )
514 local_bbox.clip(detector_bbox)
515 if local_bbox.getArea() == 0: 515 ↛ 516line 515 didn't jump to line 516 because the condition on line 515 was never true
516 return np.nan
518 model_local = ImageF(local_bbox)
519 warpImage(model_local, model_image_legacy, cutout_to_pixels, warp_control)
520 # Replace NaN values with zero; prevents NaN corners after rotation.
521 invalid_model = ~np.isfinite(model_local.array)
522 if np.any(invalid_model): 522 ↛ 525line 522 didn't jump to line 525 because the condition on line 522 was always true
523 model_local.array[invalid_model] = 0.0
525 local_mi = exposure.maskedImage[local_bbox]
526 data = local_mi.image.array.astype(np.float64)
527 variance = local_mi.variance.array.astype(np.float64)
528 model = model_local.array.astype(np.float64)
529 grid_y, grid_x = np.mgrid[
530 local_bbox.getMinY() : local_bbox.getMaxY() + 1,
531 local_bbox.getMinX() : local_bbox.getMaxX() + 1,
532 ]
533 background_terms = np.vstack(
534 [(grid_x**i * grid_y**j).astype(np.float64).ravel() for i, j in self._bg_powers]
535 ).T
537 bitmask = local_mi.mask.getPlaneBitMask(self.config.bad_mask_planes)
538 good = (local_mi.mask.array & bitmask) == 0
539 good &= np.isfinite(data) & np.isfinite(variance) & np.isfinite(model)
540 good &= variance > 0
541 good &= model > self.config.min_model_value
543 if np.count_nonzero(good) < self.config.min_fit_pixels: 543 ↛ 544line 543 didn't jump to line 544 because the condition on line 543 was never true
544 return np.nan
546 good_flat = good.ravel()
547 sigma = np.sqrt(variance[good])
548 normalized_data = data[good] / sigma
549 normalized_model = model[good] / sigma
550 normalized_background_terms = background_terms[good_flat] / sigma[:, None]
551 fit_matrix = np.hstack([normalized_model[:, None], normalized_background_terms])
553 try:
554 solution, _, _, _ = np.linalg.lstsq(fit_matrix, normalized_data, rcond=None)
555 except np.linalg.LinAlgError:
556 return np.nan
558 amplitude = float(solution[0])
559 if amplitude <= 0: 559 ↛ 560line 559 didn't jump to line 560 because the condition on line 559 was never true
560 return np.nan # Signifies a failed fit; no subtraction is done.
562 local_mi.image.array -= amplitude * model
563 return amplitude