Coverage for python/lsst/meas/extensions/scarlet/utils.py: 89%
205 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-25 01:27 -0700
1import logging
2import warnings
4import lsst.geom as geom
5import lsst.scarlet.lite as scl
6import numpy as np
7from scipy.signal import convolve
8from lsst.afw.detection import InvalidPsfError, Footprint as afwFootprint
9from lsst.afw.image import (
10 IncompleteDataError,
11 MultibandExposure,
12 MultibandImage,
13 Exposure,
14)
15from lsst.afw.image.utils import projectImage
16from lsst.afw.table import SourceCatalog
17from lsst.geom import Box2I, Point2D, Point2I
18from lsst.pipe.base import NoWorkFound
20logger = logging.getLogger(__name__)
22__all__ = [
23 "defaultBadPixelMasks",
24 "scarletBoxToBBox",
25 "bboxToScarletBox",
26 "nonzeroBandSupport",
27 "multiband_convolve",
28 "computePsfKernelImage",
29 "computeNearestPsf",
30 "computeNearestPsfMultiBand",
31 "buildObservation",
32 "calcChi2",
33]
35defaultBadPixelMasks = ["BAD", "NO_DATA", "SAT", "SUSPECT", "EDGE"]
38def scarletBoxToBBox(box: scl.Box, xy0: geom.Point2I = geom.Point2I()) -> geom.Box2I:
39 """Convert a scarlet_lite Box into a Box2I.
41 Parameters
42 ----------
43 box:
44 The scarlet bounding box to convert.
45 xy0:
46 An additional offset to add to the scarlet box.
47 This is common since scarlet sources have an origin of
48 `(0,0)` at the lower left corner of the blend while
49 the blend itself is likely to have an offset in the
50 `Exposure`.
52 Returns
53 -------
54 bbox:
55 The converted bounding box.
56 """
57 xy0 = geom.Point2I(box.origin[-1] + xy0.x, box.origin[-2] + xy0.y)
58 extent = geom.Extent2I(box.shape[-1], box.shape[-2])
59 return geom.Box2I(xy0, extent)
62def bboxToScarletBox(bbox: geom.Box2I, xy0: geom.Point2I = geom.Point2I()) -> scl.Box:
63 """Convert a Box2I into a scarlet_lite Box.
65 Parameters
66 ----------
67 bbox:
68 The Box2I to convert into a scarlet `Box`.
69 xy0:
70 An overall offset to subtract from the `Box2I`.
71 This is common in blends, where `xy0` is the minimum pixel
72 location of the blend and `bbox` is the box containing
73 a source in the blend.
75 Returns
76 -------
77 box:
78 A scarlet `Box` that is more useful for slicing image data
79 as a numpy array.
80 """
81 origin = (bbox.getMinY() - xy0.y, bbox.getMinX() - xy0.x)
82 return scl.Box((bbox.getHeight(), bbox.getWidth()), origin)
85def nonzeroBandSupport(data: np.ndarray) -> np.ndarray:
86 """Return the per-pixel support of a multi-band model.
88 A pixel is in the support whenever any band's value is non-zero.
89 This is the canonical "spatial extent of a model across bands"
90 test; the alternative idioms ``data > 0`` and
91 ``np.max(data, axis=0) != 0`` either exclude negative-valued
92 pixels outright or exclude pixels whose largest band value is
93 exactly zero, both of which under-count the true spatial extent.
95 Parameters
96 ----------
97 data :
98 A ``(bands, height, width)`` array of model values.
100 Returns
101 -------
102 support :
103 A ``(height, width)`` boolean mask, ``True`` at pixels where
104 at least one band is non-zero.
105 """
106 return np.any(data != 0, axis=0)
109def multiband_convolve(images: np.ndarray, psfs: np.ndarray) -> np.ndarray:
110 """Convolve a multi-band image with the PSF in each band.
112 `images` and `psfs` should have dimensions `(bands, height, width)`.
114 Parameters
115 ----------
116 images :
117 The multi-band images to convolve.
118 psfs :
119 The PSF for each band.
121 Returns
122 -------
123 result :
124 The convolved images.
125 """
126 result = np.zeros(images.shape, dtype=images.dtype)
127 for bidx, (image, psf) in enumerate(zip(images, psfs, strict=True)):
128 result[bidx] = convolve(image, psf, mode="same")
129 return result
132def computePsfKernelImage(mExposure, psfCenter, catalog=None):
133 """Compute the PSF kernel image and update the multiband exposure
134 if not all of the PSF images could be computed.
136 Parameters
137 ----------
138 psfCenter : `tuple` or `Point2I` or `Point2D`
139 The location `(x, y)` used as the center of the PSF.
140 catalog :
141 Deprecated and ignored. Retained for signature stability; will
142 be removed after v31. Passing a non-``None`` value emits a
143 ``FutureWarning``. For nearest-PSF fallback at a different
144 location, call ``computeNearestPsfMultiBand`` instead.
146 Returns
147 -------
148 psfModels : `np.ndarray`
149 The multiband PSF image
150 mExposure : `MultibandExposure`
151 The exposure, updated to only use bands that
152 successfully generated a PSF image.
153 """
154 if catalog is not None:
155 warnings.warn(
156 "The `catalog` parameter to `computePsfKernelImage` is "
157 "deprecated and ignored; it will be removed after v31. "
158 "For nearest-PSF fallback, use `computeNearestPsfMultiBand`.",
159 FutureWarning, stacklevel=2,
160 )
161 if not isinstance(psfCenter, geom.Point2D): 161 ↛ 162line 161 didn't jump to line 162 because the condition on line 161 was never true
162 psfCenter = geom.Point2D(*psfCenter)
164 try:
165 psfModels = mExposure.computePsfKernelImage(psfCenter)
166 except IncompleteDataError as e:
167 psfModels = e.partialPsf
168 if psfModels is None: 168 ↛ 171line 168 didn't jump to line 171 because the condition on line 168 was always true
169 return None, None
170 # Use only the bands that successfully generated a PSF image.
171 bands = psfModels.bands
172 mExposure = mExposure[bands,]
173 if len(bands) == 1:
174 # Only a single band generated a PSF, so the MultibandExposure
175 # became a single band ExposureF.
176 # Convert the result back into a MultibandExposure.
177 mExposure = MultibandExposure.fromExposures(bands, [mExposure])
178 return psfModels.array, mExposure
181def computeNearestPsf(
182 calexp: Exposure,
183 catalog: SourceCatalog,
184 band: str | None = None,
185 psfCenter: Point2D | None = None,
186) -> tuple[np.ndarray, Point2D, float] | tuple[None, None, None]:
187 """Create a PSF image at the nearest valid location
189 Sometimes not all locations in an image can generate a PSF image so the
190 source catalog is used to find the nearest valid location.
192 Parameters
193 ----------
194 calexp :
195 The exposure.
196 catalog :
197 The catalog.
198 band :
199 The band of the exposure used to filter the catalog by only
200 selecting sources that have a
201 If band is ``None`` then the full catalog is used.
202 psfCenter :
203 The location of the PSF image.
204 If no location is provided, the center of the exposure is used.
206 Returns
207 -------
208 psf :
209 The PSF image.
210 location :
211 The location of the PSF image.
212 diff :
213 The difference between the requested location and the
214 nearest valid location.
215 """
216 if psfCenter is None: 216 ↛ 217line 216 didn't jump to line 217 because the condition on line 216 was never true
217 psfCenter = calexp.getBBox().getCenter()
219 if not isinstance(psfCenter, geom.Point2D): 219 ↛ 220line 219 didn't jump to line 220 because the condition on line 219 was never true
220 psfCenter = geom.Point2D(*psfCenter)
222 try:
223 psf = calexp.getPsf().computeKernelImage(psfCenter)
224 return psf, psfCenter, 0
225 except InvalidPsfError:
226 pass
228 xc, yc = psfCenter
230 # Only select records that have detections in this band
231 if band is not None: 231 ↛ 232line 231 didn't jump to line 232 because the condition on line 231 was never true
232 sources = catalog[catalog[f'merge_footprint_{band}']]
233 else:
234 sources = catalog
236 # Get the peaks of all of the sources
237 x = []
238 y = []
239 for src in sources:
240 for peak in src.getFootprint().peaks:
241 if band is None or peak[f'merge_peak_{band}']: 241 ↛ 240line 241 didn't jump to line 240 because the condition on line 241 was always true
242 x.append(peak['i_x'])
243 y.append(peak['i_y'])
244 x = np.array(x)
245 y = np.array(y)
247 # Sort the peaks based on their distance to the location
248 diff_x = x - xc
249 diff_y = y - yc
250 sorted_indices = np.argsort(diff_x**2 + diff_y**2)
252 # Iterate over sources until a location is found that can generate a PSF
253 psf = None
254 for ref_index in sorted_indices:
255 try:
256 psf = calexp.getPsf().computeKernelImage(Point2D(x[ref_index], y[ref_index]))
257 break
258 except InvalidPsfError:
259 pass
260 if psf is None:
261 return None, None, None
262 newLocation = Point2D(x[ref_index], y[ref_index])
263 diff = np.sqrt(diff_x[ref_index]**2 + diff_y[ref_index]**2)
265 return psf, newLocation, diff
268def _sortedCatalogPositions(
269 catalog: SourceCatalog | None,
270 psfCenter: Point2D,
271) -> list[Point2D]:
272 """Catalog peak positions, sorted by distance from ``psfCenter``."""
273 if catalog is None: 273 ↛ 274line 273 didn't jump to line 274 because the condition on line 273 was never true
274 return []
275 xs: list[float] = []
276 ys: list[float] = []
277 for src in catalog:
278 for peak in src.getFootprint().peaks:
279 xs.append(peak["i_x"])
280 ys.append(peak["i_y"])
281 if not xs:
282 return []
283 xs_a = np.asarray(xs)
284 ys_a = np.asarray(ys)
285 xc, yc = psfCenter
286 order = np.argsort((xs_a - xc) ** 2 + (ys_a - yc) ** 2)
287 return [Point2D(float(xs_a[i]), float(ys_a[i])) for i in order]
290def computeNearestPsfMultiBand(
291 mExposure: MultibandExposure,
292 psfCenter: tuple[int, int] | geom.Point2I | geom.Point2D,
293 catalog: SourceCatalog | None,
294) -> tuple[MultibandImage | None, MultibandExposure | None]:
295 """Compute a multiband PSF kernel image at or near the requested center.
297 The PSF kernel is computed at ``psfCenter`` in every band where the
298 PSF model is valid there — this is the hot path and the only work
299 done when no fallback is needed. For any band whose PSF is invalid
300 at ``psfCenter``, the function walks ``catalog`` peak positions
301 sorted by distance from the requested center and accepts the first
302 position that works in every failing band as a common fallback.
303 If that fallback location also works in the bands that already
304 succeeded at the center, the function "upgrades" by sampling every
305 band at that one location, so the multiband PSF is genuinely at a
306 single sky point; otherwise the successful bands keep their center
307 PSF and only the failing bands use the common fallback, and a
308 warning is logged. When no single fallback works for every failing
309 band, each failing band falls back to its own nearest valid
310 catalog position (per-band fallback), and a warning is logged.
311 Bands with no valid PSF anywhere are dropped from the returned
312 multiband exposure, matching the historical incomplete-PSF
313 behavior.
315 PSF computations are not duplicated: each ``(band, position)`` pair
316 is evaluated at most once, and at function return the kept PSFs are
317 held one per band.
319 Parameters
320 ----------
321 mExposure :
322 The multi-band exposure.
323 psfCenter :
324 The location ``(x, y)`` used as the center of the PSF.
325 catalog :
326 Source catalog whose peak positions are candidate fallback
327 locations. If `None`, no fallback search is performed and a
328 band whose PSF is invalid at ``psfCenter`` is dropped.
330 Returns
331 -------
332 mPsf :
333 The multiband PSF kernel image, or `None` if no band produced a
334 valid PSF.
335 mExposure :
336 The input exposure restricted to bands that produced a valid
337 PSF.
338 """
339 if not isinstance(psfCenter, Point2D): 339 ↛ 340line 339 didn't jump to line 340 because the condition on line 339 was never true
340 psfCenter = Point2D(*psfCenter)
342 bands = tuple(mExposure.bands)
344 # Stage 1: try every band at the requested center.
345 psfs: dict = {}
346 locations: dict[str, Point2D] = {}
347 failingBands: list[str] = []
348 for band in bands:
349 try:
350 psfs[band] = mExposure[band,].getPsf().computeKernelImage(psfCenter)
351 locations[band] = psfCenter
352 except InvalidPsfError:
353 failingBands.append(band)
355 if failingBands:
356 # Stage 2: walk catalog candidates sorted by distance from the
357 # requested center and accept the first that works for every
358 # failing band. The bands that already succeeded at the center
359 # are not part of the search — their PSFs are already at a
360 # strictly better position than any fallback could provide.
361 candidates = _sortedCatalogPositions(catalog, psfCenter)
362 commonLocation: Point2D | None = None
363 commonPsfs: dict = {}
364 # The closest valid PSF found so far for each failing band,
365 # populated as we walk the candidate list. If no single
366 # candidate works for every failing band, these are the per-
367 # band fallbacks. Each ``(band, position)`` is evaluated at
368 # most once across the walk.
369 closestPsfs: dict = {}
370 closestLocations: dict = {}
371 for pos in candidates:
372 tentative: dict = {}
373 allOk = True
374 for band in failingBands:
375 try:
376 psf = (
377 mExposure[band,].getPsf().computeKernelImage(pos)
378 )
379 except InvalidPsfError:
380 # Mark the candidate as not common, but keep trying
381 # the remaining bands at this candidate so a band
382 # that *is* valid here still gets the chance to be
383 # cached at its true closest position.
384 allOk = False
385 continue
386 tentative[band] = psf
387 if band not in closestPsfs:
388 closestPsfs[band] = psf
389 closestLocations[band] = pos
390 if allOk:
391 commonLocation = pos
392 commonPsfs = tentative
393 break
395 if commonLocation is not None:
396 # Stage 2 upgrade: if the common fallback is also valid in
397 # the bands that succeeded at the center, switch every band
398 # to it so the multiband PSF is sampled at a single sky
399 # location. Otherwise, keep center PSFs for the successful
400 # bands and use the fallback only for the failing ones.
401 upgradePsfs: dict = {}
402 canUpgrade = True
403 for band in psfs:
404 try:
405 upgradePsfs[band] = (
406 mExposure[band,].getPsf().computeKernelImage(commonLocation)
407 )
408 except InvalidPsfError:
409 canUpgrade = False
410 break
412 if canUpgrade:
413 psfs = {**upgradePsfs, **commonPsfs}
414 else:
415 logger.warning(
416 "Multiband PSF falls back at two locations: bands %s at "
417 "the requested center %s; bands %s at %s.",
418 list(psfs), psfCenter, failingBands, commonLocation,
419 )
420 for band in failingBands:
421 psfs[band] = commonPsfs[band]
422 else:
423 # No common fallback location: use the per-band closest
424 # PSFs that the candidate walk already collected. Bands
425 # without any valid PSF in the walk are absent from
426 # closestPsfs and will be dropped below.
427 psfs.update(closestPsfs)
428 locations.update(closestLocations)
429 if any(b in psfs for b in failingBands):
430 logger.warning(
431 "Multiband PSF: no single fallback location works for "
432 "every band; per-band fallback locations %s.",
433 {b: locations[b] for b in bands if b in psfs},
434 )
436 if len(psfs) == 0:
437 return None, None
439 # Project each kept band's PSF onto the union bbox so the multiband
440 # image has consistent shape across bands.
441 left = np.min([psf.getBBox().getMinX() for psf in psfs.values()])
442 bottom = np.min([psf.getBBox().getMinY() for psf in psfs.values()])
443 right = np.max([psf.getBBox().getMaxX() for psf in psfs.values()])
444 top = np.max([psf.getBBox().getMaxY() for psf in psfs.values()])
445 bbox = Box2I(Point2I(left, bottom), Point2I(right, top))
447 # Ensure that the returned multiband PSF and exposure only contain
448 # the bands for which a valid PSF was found, in the same order as the
449 # input exposure's bands.
450 bandsKept = tuple(b for b in bands if b in psfs)
451 psf_images = [projectImage(psfs[b], bbox) for b in bandsKept]
452 mPsf = MultibandImage.fromImages(bandsKept, psf_images)
454 if len(bandsKept) < len(bands):
455 mExposure = mExposure[bandsKept,]
456 if len(bandsKept) == 1: 456 ↛ 459line 456 didn't jump to line 459 because the condition on line 456 was never true
457 # ``mExposure[(band,),]`` for a single band returns an
458 # ExposureF rather than a MultibandExposure; wrap it back.
459 mExposure = MultibandExposure.fromExposures(bandsKept, [mExposure])
461 return mPsf.array, mExposure
464def buildObservation(
465 modelPsf: np.ndarray,
466 psfCenter: tuple[int, int] | geom.Point2I | geom.Point2D,
467 mExposure: MultibandExposure,
468 badPixelMasks: list[str] | None = None,
469 footprint: afwFootprint = None,
470 useWeights: bool = True,
471 convolutionType: str = "real",
472 catalog: SourceCatalog | None = None,
473) -> scl.Observation:
474 """Generate an Observation from a set of arguments.
476 Make the generation and reconstruction of a scarlet model consistent
477 by building an `Observation` from a set of arguments.
479 Parameters
480 ----------
481 modelPsf :
482 The 2D model of the PSF in the partially deconvolved space.
483 psfCenter :
484 The location `(x, y)` used as the center of the PSF.
485 mExposure :
486 The multi-band exposure that the model represents.
487 If `mExposure` is `None` then no image, variance, or weights are
488 attached to the observation.
489 footprint :
490 The footprint that is being fit.
491 If `footprint` is `None` then the weights are not updated to mask
492 out pixels not contained in the footprint.
493 badPixelMasks :
494 The keys from the bit mask plane used to mask out pixels
495 during the fit.
496 If `badPixelMasks` is `None` then the default values from
497 `ScarletDeblendConfig.badMask` are used.
498 useWeights :
499 Whether or not fitting should use inverse variance weights to
500 calculate the log-likelihood.
501 convolutionType :
502 The type of convolution to use (either "real" or "fft").
503 When reconstructing an image it is advised to use "real" to avoid
504 polluting the footprint with artifacts from the fft.
505 catalog :
506 A source catalog to use for PSFs that cannot be determined at
507 the center of the image.
509 Returns
510 -------
511 observation:
512 The observation constructed from the input parameters.
513 """
514 # Initialize the observed PSFs
515 if not isinstance(psfCenter, geom.Point2D):
516 psfCenter = geom.Point2D(*psfCenter)
517 if catalog is None:
518 psfModels, mExposure = computePsfKernelImage(mExposure, psfCenter)
519 else:
520 psfModels, mExposure = computeNearestPsfMultiBand(mExposure, psfCenter, catalog)
522 if psfModels is None:
523 raise NoWorkFound("No valid PSF could be obtained for building the observation")
525 # Use the inverse variance as the weights
526 if useWeights: 526 ↛ 534line 526 didn't jump to line 534 because the condition on line 526 was always true
527 # Zero/NaN variance produces inf/NaN weights here; the next line
528 # zeros them deliberately. Silence the spurious RuntimeWarnings
529 # the division would otherwise emit on those pixels.
530 with np.errstate(divide="ignore", invalid="ignore"):
531 weights = 1 / mExposure.variance.array
532 weights[~np.isfinite(weights)] = 0
533 else:
534 weights = np.ones_like(mExposure.image.array)
536 # Mask out bad pixels
537 if badPixelMasks is None:
538 badPixelMasks = defaultBadPixelMasks
539 badPixels = mExposure.mask.getPlaneBitMask(badPixelMasks)
540 mask = mExposure.mask.array & badPixels
541 weights[mask > 0] = 0
543 if footprint is not None: 543 ↛ 545line 543 didn't jump to line 545 because the condition on line 543 was never true
544 # Mask out the pixels outside the footprint
545 weights *= footprint.spans.asArray()
547 # Mask out non-finite pixels
548 image = mExposure.image.array.copy()
549 weights[~np.isfinite(image)] = 0
550 image[~np.isfinite(image)] = 0
552 return scl.Observation(
553 images=image,
554 variance=mExposure.variance.array,
555 weights=weights,
556 psfs=psfModels,
557 model_psf=modelPsf[None, :, :],
558 convolution_mode=convolutionType,
559 bands=mExposure.bands,
560 bbox=bboxToScarletBox(mExposure.getBBox()),
561 )
564def calcChi2(
565 model: scl.Image,
566 observation: scl.Observation,
567 footprint: np.ndarray | None = None,
568 doConvolve: bool = True,
569) -> scl.Image:
570 """Calculate the chi2 image for a model.
572 Parameters
573 ----------
574 model :
575 The model used to calculate the chi2.
576 observation :
577 The observation used to calculate the chi2.
578 footprint :
579 The footprint to use when calculating the chi2.
580 If `footprint` is `None` then the footprint is calculated
581 to be the pixels where the model is greater than 0.
582 doConvolve :
583 Whether or not to convolve the model with the PSF.
585 Returns
586 -------
587 chi2 :
588 The chi2/pixel image for the model.
589 """
590 if doConvolve: 590 ↛ 592line 590 didn't jump to line 592 because the condition on line 590 was always true
591 model = observation.convolve(model)
592 if footprint is None: 592 ↛ 593line 592 didn't jump to line 593 because the condition on line 592 was never true
593 footprint = model.data > 0
594 bbox = model.bbox
595 nBands = len(observation.images.bands)
596 residual = (observation.images[:, bbox].data - model.data) * footprint
597 cuts = observation.variance[:, bbox].data != 0
598 chi2Data = np.zeros(residual.shape, dtype=residual.dtype)
599 chi2Data[cuts] = residual[cuts]**2 / observation.variance[:, bbox].data[cuts] / nBands
600 chi2 = scl.Image(
601 chi2Data,
602 bands=model.bands,
603 yx0=model.yx0,
604 )
605 return chi2