Coverage for python/lsst/pipe/tasks/extended_psf/extended_psf_stack.py: 18%
203 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-05-27 08:30 +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 "ExtendedPsfStackConnections",
24 "ExtendedPsfStackConfig",
25 "ExtendedPsfStackTask",
26)
28from collections.abc import Sequence
30import numpy as np
31from astropy.modeling.fitting import LMLSQFitter
32from astropy.modeling.models import Moffat2D
34from lsst.afw.image import MaskedImageF
35from lsst.afw.math import StatisticsControl, statisticsStack, stringToStatisticsProperty
36from lsst.images import Image
37from lsst.pex.config import Field, ListField
38from lsst.pipe.base import PipelineTask, PipelineTaskConfig, PipelineTaskConnections, Struct
39from lsst.pipe.base.connectionTypes import Input, Output
40from lsst.utils.timer import timeMethod
42from .extended_psf_candidates import ExtendedPsfCandidate, ExtendedPsfCandidates
43from .extended_psf_cutout import NEIGHBOR_MASK_PLANE
44from .extended_psf_fit import ExtendedPsfMoffatFit
45from .extended_psf_image import ExtendedPsfImage, ExtendedPsfImageInfo
48class ExtendedPsfStackConnections(
49 PipelineTaskConnections,
50 dimensions=("instrument", "visit", "detector"),
51):
52 """Connections for ExtendedPsfStackTask."""
54 extended_psf_candidates = Input(
55 name="extended_psf_candidates",
56 storageClass="ExtendedPsfCandidates",
57 doc="Set of preprocessed cutouts, each centered on a single star.",
58 dimensions=("visit", "detector"),
59 )
60 extended_psf = Output(
61 name="extended_psf",
62 storageClass="ExtendedPsfImage",
63 doc="Extended PSF model built from stacking candidate cutouts.",
64 dimensions=("visit", "detector"),
65 )
68class ExtendedPsfStackConfig(
69 PipelineTaskConfig,
70 pipelineConnections=ExtendedPsfStackConnections,
71):
72 """Configuration parameters for ExtendedPsfStackTask."""
74 # Candidate selection
75 bad_mask_planes = ListField[str](
76 doc="Mask planes that identify excluded (masked) pixels.",
77 default=[
78 "BAD",
79 "CR",
80 "CROSSTALK",
81 "EDGE",
82 "NO_DATA",
83 "SAT",
84 "SUSPECT",
85 "UNMASKEDNAN",
86 NEIGHBOR_MASK_PLANE,
87 ],
88 )
89 psf_masked_flux_frac_threshold = Field[float](
90 doc="Maximum allowed fraction of masked PSF flux for fitting to occur.",
91 default=0.97,
92 )
93 min_focal_plane_radius = Field[float](
94 doc="Minimum distance to the center of the focal plane, in mm. "
95 "Candidates with smaller focal plane radius will be omitted.",
96 default=0.0,
97 )
98 max_focal_plane_radius = Field[float](
99 doc="Maximum distance to the center of the focal plane, in mm. "
100 "Candidates with larger focal plane radius will be omitted.",
101 default=np.inf,
102 )
104 # Stacking control
105 stack_type = Field[str](
106 default="MEANCLIP",
107 doc="Statistic name to use for stacking (from `~lsst.afw.math.Property`)",
108 )
109 stack_num_sigma_clip = Field[float](
110 doc="Number of sigma to use for clipping when stacking.",
111 default=3.0,
112 )
113 stack_num_iter = Field[int](
114 doc="Number of iterations to use for clipping when stacking.",
115 default=5,
116 )
118 # Fitting
119 use_median_variance = Field[bool](
120 doc="Use the median of the variance plane for fitting.",
121 default=False,
122 )
123 fit_iterations = Field[int](
124 doc="Number of iterations over pedestal-gradient and scaling fit.",
125 default=5,
126 )
127 bg_order = Field[int](
128 doc="Order of polynomial to fit for background in candidate cutouts.",
129 default=2,
130 )
133class ExtendedPsfStackTask(PipelineTask):
134 """Stack candidate cutouts to produce an extended PSF model."""
136 ConfigClass = ExtendedPsfStackConfig
137 _DefaultName = "extendedPsfStack"
138 config: ExtendedPsfStackConfig
140 def __init__(self, initInputs=None, *args, **kwargs):
141 super().__init__(*args, **kwargs)
143 order = self.config.bg_order
144 self._bg_powers = [(i, j) for i in range(order + 1) for j in range(order + 1) if i + j <= order]
146 def runQuantum(self, butlerQC, inputRefs, outputRefs):
147 inputs = butlerQC.get(inputRefs)
148 output = self.run(**inputs)
149 # Guard against empty outputs, resulting from failed fitting runs.
150 if output:
151 butlerQC.put(output, outputRefs)
153 @timeMethod
154 def run(
155 self,
156 extended_psf_candidates: ExtendedPsfCandidates,
157 ):
158 """Run the stacking task to produce an extended PSF model.
160 Parameters
161 ----------
162 extended_psf_candidates : `ExtendedPsfCandidates`
163 Set of preprocessed cutouts, each centered on a single star.
165 Returns
166 -------
167 extended_psf :
168 `~lsst.pipe.tasks.extended_psf.ExtendedPsfImage` or `None`
169 The extended PSF model if stacking succeeds; `None` if stacking
170 was unsuccessful (e.g, no valid candidates).
171 """
172 candidates = self._get_candidates(extended_psf_candidates)
173 if len(candidates) == 0:
174 return None
176 extended_psf = self._make_extended_psf(candidates)
177 if extended_psf is None:
178 return None
180 fit = self._fit_moffat(extended_psf)
181 extended_psf = ExtendedPsfImage(
182 image=extended_psf.image,
183 variance=extended_psf.variance,
184 info=extended_psf.info,
185 fit=fit,
186 metadata=extended_psf.metadata,
187 )
189 return Struct(extended_psf=extended_psf)
191 @staticmethod
192 def _focal_plane_radius_mm(candidate: ExtendedPsfCandidate) -> float | None:
193 radius = candidate.star_info.focal_plane_radius
194 if radius is None:
195 return None
196 if hasattr(radius, "to_value"):
197 return float(radius.to_value("mm"))
198 return float(radius)
200 def _get_candidates(self, extended_psf_candidates: ExtendedPsfCandidates) -> list[ExtendedPsfCandidate]:
201 """Get candidates that are within the configured focal-plane radius.
203 Parameters
204 ----------
205 extended_psf_candidates : `ExtendedPsfCandidates`
206 Candidate cutout collection.
208 Returns
209 -------
210 candidates : `list` [`ExtendedPsfCandidate`]
211 Candidates within the configured focal-plane radius range.
212 """
213 candidates = []
214 for candidate in extended_psf_candidates:
215 radius_mm = self._focal_plane_radius_mm(candidate)
216 if radius_mm is None:
217 continue
218 if self.config.min_focal_plane_radius <= radius_mm <= self.config.max_focal_plane_radius:
219 candidates.append(candidate)
221 if len(candidates) == 0:
222 self.log.warning(
223 "No candidate cutouts found within focal-plane radius range [%s, %s] mm.",
224 self.config.min_focal_plane_radius,
225 self.config.max_focal_plane_radius,
226 )
227 elif len(candidates) < len(extended_psf_candidates):
228 self.log.info(
229 "Only %d of %d candidate cutouts are within focal-plane radius range [%s, %s] mm.",
230 len(candidates),
231 len(extended_psf_candidates),
232 self.config.min_focal_plane_radius,
233 self.config.max_focal_plane_radius,
234 )
236 self.log.info(
237 "Constructing an extended PSF from %d candidate cutout%s.",
238 len(candidates),
239 "s" if len(candidates) > 1 else "",
240 )
241 return candidates
243 def _make_extended_psf(self, candidates: list[ExtendedPsfCandidate]) -> ExtendedPsfImage | None:
244 """Stack candidate cutouts to produce an extended PSF model.
246 Parameters
247 ----------
248 candidates : `list` [`ExtendedPsfCandidate`]
249 Candidate cutouts to stack.
251 Returns
252 -------
253 extended_psf : `ExtendedPsfImage` or `None`
254 The stacked extended PSF model, or `None` if stacking fails.
255 """
256 stack_type_property = stringToStatisticsProperty(self.config.stack_type)
257 statistics_control = StatisticsControl(
258 numSigmaClip=self.config.stack_num_sigma_clip,
259 numIter=self.config.stack_num_iter,
260 )
261 bad_mask_bit_mask = (
262 candidates[0].to_legacy(copy=False).mask.getPlaneBitMask(self.config.bad_mask_planes)
263 )
264 statistics_control.setAndMask(bad_mask_bit_mask)
266 psf_amplitudes: dict[int, float | None] = {
267 candidate_id: None for candidate_id in range(len(candidates))
268 }
270 extended_psf = None
271 # Iteratively fit and normalize candidate cutouts, then stack
272 for iteration in range(self.config.fit_iterations):
273 normalized_cutouts = []
274 # Loop over all candidates
275 for candidate_id, candidate in enumerate(candidates):
276 # Use the PSF kernel image for the first iteration
277 if extended_psf is None:
278 psf_image = candidate.psf_kernel_image
279 else:
280 psf_image = extended_psf.image
281 normalized_image, psf_amplitude = self._fit_and_normalize(
282 candidate,
283 psf_image,
284 psf_amplitudes[candidate_id],
285 )
286 if normalized_image is not None and psf_amplitude is not None:
287 normalized_cutouts.append(normalized_image)
288 psf_amplitudes[candidate_id] = psf_amplitude
290 if len(normalized_cutouts) == 0:
291 self.log.warning(
292 "Iteration %d: no candidates were successfully fit and normalized. "
293 "No further iterations will be attempted.",
294 iteration + 1,
295 )
296 break
298 self.log.info(
299 "Iteration %d: fit and normalized %d out of %d candidate cutout%s using the %s model.",
300 iteration + 1,
301 len(normalized_cutouts),
302 len(candidates),
303 "s" if len(candidates) > 1 else "",
304 "kernel-image" if extended_psf is None else "extended PSF",
305 )
307 # statisticsStack requires legacy MaskedImageF.
308 normalized_cutouts_legacy = [cutout.to_legacy(copy=False) for cutout in normalized_cutouts]
309 extended_psf_legacy = MaskedImageF(normalized_cutouts_legacy[0].getBBox())
310 statisticsStack(
311 extended_psf_legacy,
312 normalized_cutouts_legacy,
313 stack_type_property,
314 statistics_control,
315 )
316 extended_psf = ExtendedPsfImage(
317 image=Image.from_legacy(extended_psf_legacy.image),
318 variance=Image.from_legacy(extended_psf_legacy.variance),
319 info=ExtendedPsfImageInfo(n_stars=len(candidates)),
320 )
322 return extended_psf
324 def _fit_and_normalize(
325 self,
326 candidate: ExtendedPsfCandidate,
327 psf_image: Image,
328 psf_amplitude: float | None,
329 ) -> tuple[ExtendedPsfCandidate | None, float | None]:
330 """Fit and normalize a single candidate cutout.
332 Parameters
333 ----------
334 candidate : `ExtendedPsfCandidate`
335 Candidate cutout image to fit.
336 psf_image : `~lsst.images.Image`
337 PSF model image for fitting.
338 psf_amplitude : `float` | `None`
339 Prior PSF scaling to use during background fitting. If `None`, no
340 PSF subtraction is performed during background fitting.
342 Returns
343 -------
344 candidate_normalized : `ExtendedPsfCandidate` | `None`
345 Normalized candidate cutout image, or `None` if fitting fails.
346 psf_amplitude : `float` | `None`
347 Fitted PSF amplitude, or `None` if fitting fails.
348 """
349 bit_mask_bad = candidate.mask.schema.bitmask(*self.config.bad_mask_planes)
351 # Calculate the fraction of PSF flux masked by bad pixels.
352 psf_mask = ((candidate.mask[psf_image.bbox].array & bit_mask_bad) != 0).any(axis=-1)
353 psf_masked_flux = np.dot(psf_image.array.flat, psf_mask.astype(np.float64).flat).astype(np.float64)
354 psf_masked_flux_frac = psf_masked_flux / psf_image.array.sum()
355 if psf_masked_flux_frac > self.config.psf_masked_flux_frac_threshold:
356 return None, None
358 # Fit the background.
359 bg_image = self._fit_bg(
360 candidate,
361 list(self.config.bad_mask_planes) + ["DETECTED"],
362 psf_image,
363 psf_amplitude,
364 )
365 if bg_image is None:
366 return None, None
367 candidate_fit: ExtendedPsfCandidate = candidate.copy()
368 candidate_fit.image.array -= bg_image.array
370 # Fit PSF amplitude.
371 psf_image_padded = Image(0.0, dtype=psf_image.array.dtype, bbox=candidate.bbox, unit=psf_image.unit)
372 psf_image_padded[psf_image.bbox] = psf_image
373 psf_amplitude = self._fit_psf_amplitude(
374 candidate_fit,
375 psf_image_padded,
376 self.config.bad_mask_planes,
377 )
378 if psf_amplitude is None or psf_amplitude <= 0:
379 return None, None
380 candidate_fit.image.array /= psf_amplitude
381 candidate_fit.variance.array /= psf_amplitude**2
383 return candidate_fit, psf_amplitude
385 def _fit_bg(
386 self,
387 candidate: ExtendedPsfCandidate,
388 mask_planes: Sequence[str],
389 psf_image: Image,
390 psf_amplitude: float | None,
391 ) -> Image | None:
392 """Fit a polynomial background to a candidate cutout.
394 Parameters
395 ----------
396 candidate : `ExtendedPsfCandidate`
397 Candidate cutout image to fit.
398 mask_planes : `Sequence` [`str`]
399 Mask planes for identifying bad pixels.
400 psf_image : `~lsst.images.Image`
401 PSF image to optionally subtract prior to fitting.
402 psf_amplitude : `float` | `None`
403 Amplitude to scale ``psf_image`` before subtraction. If `None`, no
404 PSF subtraction is performed.
406 Returns
407 -------
408 bg_image : `~lsst.images.Image` | `None`
409 Fitted background image, or `None` if fitting fails.
410 """
411 if psf_amplitude is not None:
412 psf_image_scaled = psf_image.copy()
413 psf_image_scaled.array *= psf_amplitude
415 good_pixels, image, sigma = self._get_span_data(
416 candidate,
417 mask_planes,
418 subtract_image=psf_image_scaled if psf_amplitude is not None else None,
419 )
420 bbox = candidate.bbox
421 grid = bbox.meshgrid()
422 grid_x, grid_y = grid.x, grid.y
423 spans_x = grid_x[good_pixels]
424 spans_y = grid_y[good_pixels]
426 design_matrix = np.ones((len(image), len(self._bg_powers)), dtype=float)
427 for k, (i, j) in enumerate(self._bg_powers):
428 if i == j == 0:
429 design_matrix[:, k] /= sigma
430 else:
431 design_matrix[:, k] = (spans_x**i * spans_y**j) / sigma
433 solutions, _ = self._solve_design_matrix(design_matrix, image)
434 if solutions is None:
435 return None
437 bg_array = sum(s * (grid_x**i * grid_y**j) for s, (i, j) in zip(solutions, self._bg_powers))
439 return Image(bg_array.astype(np.float32), bbox=candidate.bbox, unit=candidate.unit)
441 def _fit_psf_amplitude(
442 self,
443 candidate: ExtendedPsfCandidate,
444 psf_image: Image,
445 mask_planes: Sequence[str],
446 mask_zeros: bool = True,
447 ) -> float | None:
448 """Fit the amplitude of the PSF model to one candidate cutout.
450 Parameters
451 ----------
452 candidate : `ExtendedPsfCandidate`
453 Candidate cutout image to fit.
454 psf_image : `~lsst.images.Image`
455 PSF image to fit to the data.
456 mask_planes : `Sequence` [`str`]
457 Mask planes for identifying bad pixels.
458 mask_zeros : `bool`, optional
459 Whether to additionally mask zero-valued PSF pixels.
461 Returns
462 -------
463 psf_amplitude : `float` | `None`
464 Fitted PSF amplitude, or `None` if fitting fails.
466 Notes
467 -----
468 NaN and negative values in the PSF image are always masked.
469 """
470 psf_mask_array = np.isnan(psf_image.array) | (psf_image.array < 0)
471 if mask_zeros:
472 psf_mask_array |= psf_image.array == 0
474 good_pixels, image, sigma = self._get_span_data(candidate, mask_planes, psf_mask_array)
475 psf = psf_image.array[good_pixels]
476 psf /= sigma
477 design_matrix = psf[:, None]
479 solutions, _ = self._solve_design_matrix(design_matrix, image)
480 if solutions is None:
481 return None
483 return float(solutions[0])
485 def _get_span_data(
486 self,
487 candidate: ExtendedPsfCandidate,
488 mask_planes: Sequence[str],
489 additional_mask: np.ndarray | None = None,
490 subtract_image: Image | None = None,
491 ) -> tuple[np.ndarray, np.ndarray, np.ndarray | float]:
492 """Get fitting data and a boolean mask of good pixels.
494 Parameters
495 ----------
496 candidate : `ExtendedPsfCandidate`
497 Candidate cutout from which to extract fitting data.
498 mask_planes : `Sequence` [`str`]
499 Mask planes for identifying bad pixels.
500 additional_mask : `numpy.ndarray` or `None`, optional
501 Additional boolean mask to combine with the image mask.
502 subtract_image : `~lsst.images.Image` or `None`, optional
503 Image to subtract before fitting (e.g., PSF model).
505 Returns
506 -------
507 good_pixels : `numpy.ndarray`
508 Boolean array selecting good pixels.
509 image : `numpy.ndarray`
510 Pixel values at good pixels after optional subtraction.
511 sigma : `numpy.ndarray` | `float`
512 Pixel-wise uncertainties at good pixels.
513 """
514 bit_mask = candidate.mask.schema.bitmask(*mask_planes)
515 bad_pixels = ((candidate.mask.array & bit_mask) != 0).any(axis=-1)
516 if additional_mask is not None:
517 bad_pixels |= additional_mask
518 if subtract_image is None:
519 image_data = candidate.image.array
520 else:
521 image_data_image = candidate.image.copy()
522 image_data_image[subtract_image.bbox].array -= subtract_image.array
523 image_data = image_data_image.array
524 good_pixels = ~bad_pixels
525 image = image_data[good_pixels]
526 variance = candidate.variance.array[good_pixels]
527 if self.config.use_median_variance:
528 variance = float(np.median(variance))
529 if np.isscalar(variance):
530 sigma = float(np.sqrt(variance))
531 if sigma == 0:
532 sigma = np.inf
533 else:
534 sigma = np.sqrt(variance)
535 sigma[sigma == 0] = np.inf
536 return good_pixels, image / sigma, sigma
538 def _solve_design_matrix(
539 self,
540 design_matrix: np.ndarray,
541 data: np.ndarray,
542 calculate_covariance: bool = False,
543 ) -> tuple[np.ndarray | None, np.ndarray | None]:
544 """Solve a linear system for the given design matrix and data.
546 Parameters
547 ----------
548 design_matrix : `numpy.ndarray`
549 Design matrix for the linear system.
550 data : `numpy.ndarray`
551 Data vector.
552 calculate_covariance : `bool`, optional
553 Whether to return the covariance matrix.
555 Returns
556 -------
557 solutions : `numpy.ndarray` or `None`
558 Solution vector, or `None` if solving fails.
559 covariance_matrix : `numpy.ndarray` or `None`
560 Covariance matrix, if requested and available.
561 """
562 covariance_matrix = None
563 try:
564 solutions, sum_squared_residuals, *_ = np.linalg.lstsq(design_matrix, data, rcond=None)
565 if calculate_covariance:
566 covariance_matrix = np.linalg.pinv(design_matrix.T @ design_matrix)
567 except np.linalg.LinAlgError:
568 return None, None
569 if sum_squared_residuals.size == 0:
570 return None, None
571 return solutions, covariance_matrix
573 def _fit_moffat(self, extended_psf: ExtendedPsfImage, fix_center: bool = False) -> ExtendedPsfMoffatFit:
574 """Fit a Moffat2D model to the stacked extended PSF image.
576 Parameters
577 ----------
578 extended_psf : `ExtendedPsfImage`
579 Stacked extended PSF image to fit.
580 fix_center : `bool`, optional
581 Whether to fix the Moffat2D center at the image center.
583 Returns
584 -------
585 fit : `ExtendedPsfMoffatFit`
586 Moffat fit summary and fit statistics.
587 """
588 bbox = extended_psf.bbox
589 grid = bbox.meshgrid()
590 grid_x, grid_y = grid.x, grid.y
592 extended_psf_image = extended_psf.image.array
593 extended_psf_sigma = np.sqrt(extended_psf.variance.array)
595 fitter = LMLSQFitter()
596 moffat_init = Moffat2D(
597 amplitude=extended_psf_image.max(),
598 x_0=bbox.x.center,
599 y_0=bbox.y.center,
600 gamma=3.0,
601 alpha=2.5,
602 fixed={"x_0": fix_center, "y_0": fix_center},
603 )
604 weights = 1.0 / np.clip(extended_psf_sigma, 1e-12, None)
605 moffat_fit = fitter(moffat_init, grid_x, grid_y, extended_psf_image, weights=weights)
607 residuals = extended_psf_image - moffat_fit(grid_x, grid_y)
608 chi2 = np.sum((residuals / extended_psf_sigma) ** 2)
609 dof = extended_psf_image.size - len(moffat_fit.parameters)
610 reduced_chi2 = np.nan if dof <= 0 else chi2 / dof
612 self.log.info(
613 "Extended PSF Moffat fit x_0=%.2f, y_0=%.2f, gamma=%.2f, alpha=%.2f, reduced_chi2=%.2f.",
614 moffat_fit.x_0.value,
615 moffat_fit.y_0.value,
616 moffat_fit.gamma.value,
617 moffat_fit.alpha.value,
618 reduced_chi2,
619 )
621 return ExtendedPsfMoffatFit(
622 chi2=float(chi2),
623 dof=int(dof),
624 reduced_chi2=float(reduced_chi2),
625 amplitude=float(moffat_fit.amplitude.value),
626 x_0=float(moffat_fit.x_0.value),
627 y_0=float(moffat_fit.y_0.value),
628 gamma=float(moffat_fit.gamma.value),
629 alpha=float(moffat_fit.alpha.value),
630 )