Coverage for python/lsst/images/cells/_psf.py: 40%
168 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-03 01:34 +0000
1# This file is part of lsst-images.
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# Use of this source code is governed by a 3-clause BSD-style
10# license that can be found in the LICENSE file.
12from __future__ import annotations
14__all__ = ("CellPointSpreadFunction", "CellPointSpreadFunctionSerializationModel")
16from functools import cached_property
17from typing import TYPE_CHECKING, Any, ClassVar, Literal, overload
19import numpy as np
20import pydantic
22from .._cell_grid import CellGrid, CellGridBounds, CellIJ
23from .._geom import XY, YX, Bounds, BoundsError, Box, Interval
24from .._image import Image
25from ..psfs import PointSpreadFunction
26from ..serialization import (
27 ArchiveTree,
28 ArrayReferenceModel,
29 InlineArrayModel,
30 InputArchive,
31 InvalidParameterError,
32 OutputArchive,
33)
34from ..utils import round_half_up
36if TYPE_CHECKING:
37 try:
38 from lsst.afw.detection import Psf as LegacyPsf
39 from lsst.cell_coadds import StitchedPsf as LegacyStitchedPsf
40 except ImportError:
41 type LegacyPsf = Any # type: ignore[no-redef]
42 type LegacyStitchedPsf = Any # type: ignore[no-redef]
45class CellPointSpreadFunction(PointSpreadFunction):
46 """A PSF model that is at least approximately constant over cells.
48 Parameters
49 ----------
50 array
51 A 4-d array of PSF kernel images with with shape
52 ``(n_cells_y, n_cells_x, psf_shape_y, psf_shape_x)``.
53 bounds
54 Description of the cell grid and any missing cells. Array entries for
55 missing cells should be NaN.
56 resampling_kernel
57 Name of the resampling kernel to use when shifting the kernel image
58 into the stellar image.
60 Notes
61 -----
62 Unlike most PSF model types, `CellPointSpreadFunction` can be subset via
63 slicing:
65 - a bounding `.Box` for a subimage, which returns a new PSF with only the
66 cells that cover that subimage;
67 - a `CellIJ` index, which returns the kernel image for that cell.
68 -
69 """
71 def __init__(
72 self,
73 array: np.ndarray,
74 bounds: CellGridBounds,
75 resampling_kernel: Literal["lanczos3", "lanczos5"] = "lanczos5",
76 ) -> None:
77 self._array = array
78 self._bounds: CellGridBounds = bounds
79 self._resampling_kernel = resampling_kernel
81 @property
82 def grid(self) -> CellGrid:
83 """The grid that defines the PSF's cells (`CellGrid`).
85 Notes
86 -----
87 This is usually (but is not guaranteed to be) the grid for a full
88 patch, even when the PSF only covers a subimage.
89 """
90 return self._bounds.grid
92 @property
93 def bounds(self) -> CellGridBounds:
94 """The bounds where the PSF can be evaluated (`CellGridBounds`)."""
95 return self._bounds
97 @cached_property
98 def kernel_bbox(self) -> Box:
99 sy, sx = self._array.shape[2:]
100 ry = sy // 2
101 rx = sx // 2
102 return Box.factory[-ry : ry + 1, -rx : rx + 1]
104 @overload
105 def __getitem__(self, bbox: Box) -> CellPointSpreadFunction: ... 105 ↛ exitline 105 didn't return from function '__getitem__' because
106 @overload
107 def __getitem__(self, index: CellIJ) -> Image: ... 107 ↛ exitline 107 didn't return from function '__getitem__' because
109 def __getitem__(self, key: Box | CellIJ) -> CellPointSpreadFunction | Image:
110 match key:
111 case CellIJ():
112 if key in self._bounds.missing: 112 ↛ 113line 112 didn't jump to line 113 because the condition on line 112 was never true
113 raise BoundsError(f"Cell {key} is missing for this PSF.")
114 if not self._bounds.contains_cell(key): 114 ↛ 116line 114 didn't jump to line 116 because the condition on line 114 was always true
115 raise BoundsError(f"Cell {key} is out of bounds for this PSF.")
116 index = key - self._bounds.subgrid_start
117 try:
118 return Image(self._array[index.i, index.j], bbox=self.kernel_bbox)
119 except IndexError:
120 raise BoundsError(f"Cell {key} is out of bounds for this PSF.")
121 case Box(): 121 ↛ 124line 121 didn't jump to line 124 because the pattern on line 121 always matched
122 bounds, slices = self._subset_impl(self._bounds, key)
123 return CellPointSpreadFunction(self._array[slices.y, slices.x, ...].copy(), bounds=bounds)
124 case _:
125 raise TypeError("Invalid argument for CellPointSpreadFunction.__getitem__.")
127 def compute_kernel_image(self, *, x: float, y: float) -> Image:
128 index = self.grid.index_of(x=round(x), y=round(y))
129 try:
130 return self[index]
131 except Exception as err:
132 err.add_note(f"Evaluating cell PSF at x={x}, y={y}.")
133 raise
135 def compute_stellar_image(self, *, x: float, y: float) -> Image:
136 try:
137 from lsst.afw.math import offsetImage
138 from lsst.geom import Point2I
139 except ImportError as err:
140 err.add_note("CellPointSpreadFunction.compute_stellar_image cannot be used without lsst.afw.")
141 raise
142 ix = round_half_up(x)
143 dx = x - ix
144 iy = round_half_up(y)
145 dy = y - iy
146 kernel_image = self.compute_kernel_image(x=x, y=y)
147 if dx != 0 or dy != 0:
148 legacy_result = offsetImage(kernel_image.to_legacy(), dx, dy, self._resampling_kernel, 5)
149 else:
150 # This branch is equal to the other up to round-off error, but it's
151 # convenient nonetheless because it maintains exact compatibility
152 # with the legacy implementation, where the caching mechanism
153 # causes the offsetImage call to be skipped.
154 legacy_result = kernel_image.to_legacy()
155 legacy_result.setXY0(Point2I(legacy_result.getX0() + ix, legacy_result.getY0() + iy))
156 return Image.from_legacy(legacy_result)
158 def compute_stellar_bbox(self, *, x: float, y: float) -> Box:
159 # This is obviously inefficient, but it's what afw does, and hence the
160 # only easy way we've got to replicate what afw does.
161 return self.compute_stellar_image(x=x, y=y).bbox
163 def serialize(self, archive: OutputArchive[Any]) -> CellPointSpreadFunctionSerializationModel:
164 array_model = archive.add_array(self._array)
165 return CellPointSpreadFunctionSerializationModel(array=array_model, bounds=self.bounds)
167 @classmethod
168 def from_legacy(
169 cls, legacy_psf: LegacyStitchedPsf, bounds: Bounds | None = None
170 ) -> CellPointSpreadFunction:
171 # 'bounds' is accepted as an argument only for base-class
172 # compatibility; we always generate our own bounds.
173 from lsst.geom import Box2I
175 grid = CellGrid.from_legacy(legacy_psf.grid)
176 # Start with bounds that cover the entire grid.
177 bounds = CellGridBounds(grid=grid, bbox=grid.bbox)
178 # Shrink bounds to just the bbox where we have data.
179 legacy_bbox = Box2I()
180 for legacy_index in legacy_psf.images.keys():
181 legacy_bbox.include(legacy_psf.grid.bbox_of(legacy_index))
182 bounds = bounds[Box.from_legacy(legacy_bbox)]
183 # Allocate and populate the array.
184 psf_image_size_y, psf_image_size_x = legacy_psf.images.arbitrary.array.shape
185 array = np.zeros(
186 (bounds.subgrid_size.i, bounds.subgrid_size.j, psf_image_size_y, psf_image_size_x),
187 dtype=np.float64,
188 )
189 missing: set[CellIJ] = set()
190 for cell_index in bounds.cell_indices():
191 legacy_index = cell_index.to_legacy()
192 array_index = cell_index - bounds.subgrid_start
193 if legacy_index in legacy_psf.images:
194 array[array_index.i, array_index.j] = legacy_psf.images[legacy_index].array
195 else:
196 array[array_index.i, array_index.j] = np.nan
197 missing.add(cell_index)
198 # Modify the bounds one last time to account for missing cells.
199 bounds = CellGridBounds(grid=grid, bbox=bounds.bbox, missing=frozenset(missing))
200 return cls(array, bounds=bounds)
202 def to_legacy(self) -> LegacyStitchedPsf:
203 """Convert to a legacy `lsst.cell_coadds.StitchedPsf` object."""
204 from lsst.afw.image import ImageD as LegacyImageD
205 from lsst.cell_coadds import GridContainer as LegacyGridContainer
206 from lsst.cell_coadds import StitchedPsf as LegacyStitchedPsf
208 grid = self.grid.to_legacy()
209 gc = LegacyGridContainer[LegacyImageD](grid.shape)
210 for cell_index in self.bounds.cell_indices():
211 gc[cell_index.to_legacy()] = self[cell_index].to_legacy()
212 return LegacyStitchedPsf(gc, grid)
214 @staticmethod
215 def approximate_from_legacy(legacy: LegacyPsf, bounds: CellGridBounds) -> CellPointSpreadFunction:
216 """Convert from a legacy PSF by evaluating it at cell centers and
217 assuming variation within a cell is approximately constant.
219 When the legacy PSF cannot be evaluated at a cell's center (e.g. the
220 center falls outside every contributing exposure), it is instead
221 sampled at the cell's four corners and the valid samples are averaged.
223 Parameters
224 ----------
225 legacy
226 Legacy point-spread function to approximate.
227 bounds
228 Cell grid and known-missing cells.
229 """
230 from lsst.afw.detection import InvalidPsfError
232 stamps: dict[CellIJ, list[Image]] = {}
233 missing: set[CellIJ] = set()
234 x_intervals: list[Interval] = []
235 y_intervals: list[Interval] = []
237 def append_stamp_at(cell_index: CellIJ, point: XY[float] | YX[int]) -> bool:
238 try:
239 stamp = Image.from_legacy(legacy.computeKernelImage(point.to_legacy_float_point()))
240 stamps.setdefault(cell_index, []).append(stamp)
241 x_intervals.append(stamp.bbox.x)
242 y_intervals.append(stamp.bbox.y)
243 except InvalidPsfError:
244 return False
245 return True
247 for cell_index in bounds.cell_indices():
248 cell_bbox = bounds.grid.bbox_of(cell_index)
249 if not append_stamp_at(cell_index, cell_bbox.center):
250 # Average evaluations at corners instead (or at least the ones
251 # where we can evaluate). We use `sum` instead of `any` to
252 # avoid short-circuiting.
253 if not sum(append_stamp_at(cell_index, corner) for corner in cell_bbox.boundary()):
254 missing.add(cell_index)
255 continue
256 if not stamps:
257 raise ValueError("Given PSF could not be evaluated at any cell center or corner.")
258 kernel_bbox = Box(y=Interval.hull(*y_intervals), x=Interval.hull(*x_intervals))
259 array = np.full(bounds.subgrid_size.as_tuple() + kernel_bbox.shape, np.nan, dtype=np.float64)
260 for cell_index, stamps_for_cell in stamps.items():
261 assert stamps_for_cell, "Should be guaranteed by append_stamp_at."
262 array_index = cell_index - bounds.subgrid_start
263 array[array_index.i, array_index.j, :, :] = 0.0
264 for stamp in stamps_for_cell:
265 slices = stamp.bbox.slice_within(kernel_bbox)
266 array[array_index.i, array_index.j, slices.y, slices.x] += stamp.array
267 array[array_index.i, array_index.j, :, :] /= len(stamps_for_cell)
268 if missing:
269 bounds = CellGridBounds(
270 grid=bounds.grid, bbox=bounds.bbox, missing=frozenset(bounds.missing | missing)
271 )
272 return CellPointSpreadFunction(array, bounds=bounds)
274 @staticmethod
275 def _subset_impl(bounds: CellGridBounds, bbox: Box) -> tuple[CellGridBounds, YX[slice]]:
276 subset_bounds = bounds[bbox]
277 start = subset_bounds.subgrid_start - bounds.subgrid_start
278 stop = subset_bounds.subgrid_stop - bounds.subgrid_start
279 return subset_bounds, YX(y=slice(start.i, stop.i), x=slice(start.j, stop.j))
282class CellPointSpreadFunctionSerializationModel(ArchiveTree):
283 """Model used to serialize CellPointSpreadFunction objects."""
285 SCHEMA_NAME: ClassVar[str] = "cell_psf"
286 SCHEMA_VERSION: ClassVar[str] = "1.0.0"
287 MIN_READ_VERSION: ClassVar[int] = 1
288 PUBLIC_TYPE: ClassVar[type] = CellPointSpreadFunction
290 array: ArrayReferenceModel | InlineArrayModel = pydantic.Field(
291 description=(
292 "A 4-d array of PSF kernel images with with shape "
293 "(n_cells_y, n_cells_x, psf_shape_y, psf_shape_x)."
294 )
295 )
296 bounds: CellGridBounds = pydantic.Field(
297 description=(
298 "Description of the cell grid and any missing cells. Array entries for "
299 "missing cells should be NaN."
300 )
301 )
303 def deserialize(
304 self, archive: InputArchive[Any], *, bbox: Box | None = None, **kwargs: Any
305 ) -> CellPointSpreadFunction:
306 if kwargs: 306 ↛ 307line 306 didn't jump to line 307 because the condition on line 306 was never true
307 raise InvalidParameterError(
308 f"Unrecognized parameters for CellPointSpreadFunction: {set(kwargs.keys())}."
309 )
310 bounds = self.bounds
311 if bbox is not None:
312 bounds, slices = CellPointSpreadFunction._subset_impl(bounds, bbox)
313 array = archive.get_array(self.array, slices=slices)
314 else:
315 array = archive.get_array(self.array)
316 return CellPointSpreadFunction(array, bounds)