Coverage for python/lsst/images/cells/_psf.py: 43%
129 statements
« prev ^ index » next coverage.py v7.14.2, created at 2026-06-22 08:45 +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 YX, Bounds, BoundsError, Box
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.cell_coadds import StitchedPsf as LegacyStitchedPsf
39 except ImportError:
40 type LegacyStitchedPsf = Any # type: ignore[no-redef]
43class CellPointSpreadFunction(PointSpreadFunction):
44 """A PSF model that is at least approximately constant over cells.
46 Parameters
47 ----------
48 array
49 A 4-d array of PSF kernel images with with shape
50 ``(n_cells_y, n_cells_x, psf_shape_y, psf_shape_x)``.
51 bounds
52 Description of the cell grid and any missing cells. Array entries for
53 missing cells should be NaN.
54 resampling_kernel
55 Name of the resampling kernel to use when shifting the kernel image
56 into the stellar image.
58 Notes
59 -----
60 Unlike most PSF model types, `CellPointSpreadFunction` can be subset via
61 slicing:
63 - a bounding `.Box` for a subimage, which returns a new PSF with only the
64 cells that cover that subimage;
65 - a `CellIJ` index, which returns the kernel image for that cell.
66 -
67 """
69 def __init__(
70 self,
71 array: np.ndarray,
72 bounds: CellGridBounds,
73 resampling_kernel: Literal["lanczos3", "lanczos5"] = "lanczos5",
74 ):
75 self._array = array
76 self._bounds: CellGridBounds = bounds
77 self._resampling_kernel = resampling_kernel
79 @property
80 def grid(self) -> CellGrid:
81 """The grid that defines the PSF's cells (`CellGrid`).
83 Notes
84 -----
85 This is usually (but is not guaranteed to be) the grid for a full
86 patch, even when the PSF only covers a subimage.
87 """
88 return self._bounds.grid
90 @property
91 def bounds(self) -> CellGridBounds:
92 """The bounds where the PSF can be evaluated (`CellGridBounds`)."""
93 return self._bounds
95 @cached_property
96 def kernel_bbox(self) -> Box:
97 sy, sx = self._array.shape[2:]
98 ry = sy // 2
99 rx = sx // 2
100 return Box.factory[-ry : ry + 1, -rx : rx + 1]
102 @overload
103 def __getitem__(self, bbox: Box) -> CellPointSpreadFunction: ... 103 ↛ exitline 103 didn't return from function '__getitem__' because
104 @overload
105 def __getitem__(self, index: CellIJ) -> Image: ... 105 ↛ exitline 105 didn't return from function '__getitem__' because
107 def __getitem__(self, key: Box | CellIJ) -> CellPointSpreadFunction | Image:
108 match key:
109 case CellIJ(): 109 ↛ 110line 109 didn't jump to line 110 because the pattern on line 109 never matched
110 if key in self._bounds.missing:
111 raise BoundsError(f"Cell {key} is missing for this PSF.")
112 index = key - self._bounds.subgrid_start
113 try:
114 return Image(self._array[index.i, index.j], bbox=self.kernel_bbox)
115 except IndexError:
116 raise BoundsError(f"Cell {key} is out of bounds for this PSF.")
117 case Box(): 117 ↛ 120line 117 didn't jump to line 120 because the pattern on line 117 always matched
118 bounds, slices = self._subset_impl(self._bounds, key)
119 return CellPointSpreadFunction(self._array[slices.y, slices.x, ...].copy(), bounds=bounds)
120 case _:
121 raise TypeError("Invalid argument for CellPointSpreadFunction.__getitem__.")
123 def compute_kernel_image(self, *, x: float, y: float) -> Image:
124 index = self.grid.index_of(x=round(x), y=round(y))
125 try:
126 return self[index]
127 except Exception as err:
128 err.add_note(f"Evaluating cell PSF at x={x}, y={y}.")
129 raise
131 def compute_stellar_image(self, *, x: float, y: float) -> Image:
132 try:
133 from lsst.afw.math import offsetImage
134 from lsst.geom import Point2I
135 except ImportError as err:
136 err.add_note("CellPointSpreadFunction.compute_stellar_image cannot be used without lsst.afw.")
137 raise
138 ix = round_half_up(x)
139 dx = x - ix
140 iy = round_half_up(y)
141 dy = y - iy
142 kernel_image = self.compute_kernel_image(x=x, y=y)
143 if dx != 0 or dy != 0:
144 legacy_result = offsetImage(kernel_image.to_legacy(), dx, dy, self._resampling_kernel, 5)
145 else:
146 # This branch is equal to the other up to round-off error, but it's
147 # convenient nonetheless because it maintains exact compatibility
148 # with the legacy implementation, where the caching mechanism
149 # causes the offsetImage call to be skipped.
150 legacy_result = kernel_image.to_legacy()
151 legacy_result.setXY0(Point2I(legacy_result.getX0() + ix, legacy_result.getY0() + iy))
152 return Image.from_legacy(legacy_result)
154 def compute_stellar_bbox(self, *, x: float, y: float) -> Box:
155 # This is obviously inefficient, but it's what afw does, and hence the
156 # only easy way we've got to replicate what afw does.
157 return self.compute_stellar_image(x=x, y=y).bbox
159 def serialize(self, archive: OutputArchive[Any]) -> CellPointSpreadFunctionSerializationModel:
160 array_model = archive.add_array(self._array)
161 return CellPointSpreadFunctionSerializationModel(array=array_model, bounds=self.bounds)
163 @classmethod
164 def from_legacy(
165 cls, legacy_psf: LegacyStitchedPsf, bounds: Bounds | None = None
166 ) -> CellPointSpreadFunction:
167 # 'bounds' is accepted as an argument only for base-class
168 # compatibility; we always generate our own bounds.
169 from lsst.geom import Box2I
171 grid = CellGrid.from_legacy(legacy_psf.grid)
172 # Start with bounds that cover the entire grid.
173 bounds = CellGridBounds(grid=grid, bbox=grid.bbox)
174 # Shrink bounds to just the bbox where we have data.
175 legacy_bbox = Box2I()
176 for legacy_index in legacy_psf.images.keys():
177 legacy_bbox.include(legacy_psf.grid.bbox_of(legacy_index))
178 bounds = bounds[Box.from_legacy(legacy_bbox)]
179 # Allocate and populate the array.
180 psf_image_size_y, psf_image_size_x = legacy_psf.images.arbitrary.array.shape
181 array = np.zeros(
182 (bounds.subgrid_size.i, bounds.subgrid_size.j, psf_image_size_y, psf_image_size_x),
183 dtype=np.float64,
184 )
185 missing: set[CellIJ] = set()
186 for cell_index in bounds.cell_indices():
187 legacy_index = cell_index.to_legacy()
188 array_index = cell_index - bounds.subgrid_start
189 if legacy_index in legacy_psf.images:
190 array[array_index.i, array_index.j] = legacy_psf.images[legacy_index].array
191 else:
192 array[array_index.i, array_index.j] = np.nan
193 missing.add(cell_index)
194 # Modify the bounds one last time to account for missing cells.
195 bounds = CellGridBounds(grid=grid, bbox=bounds.bbox, missing=frozenset(missing))
196 return cls(array, bounds=bounds)
198 def to_legacy(self) -> LegacyStitchedPsf:
199 """Convert to a legacy `lsst.cell_coadds.StitchedPsf` object."""
200 from lsst.afw.image import ImageD as LegacyImageD
201 from lsst.cell_coadds import GridContainer as LegacyGridContainer
202 from lsst.cell_coadds import StitchedPsf as LegacyStitchedPsf
204 grid = self.grid.to_legacy()
205 gc = LegacyGridContainer[LegacyImageD](grid.shape)
206 for cell_index in self.bounds.cell_indices():
207 gc[cell_index.to_legacy()] = self[cell_index].to_legacy()
208 return LegacyStitchedPsf(gc, grid)
210 @staticmethod
211 def _subset_impl(bounds: CellGridBounds, bbox: Box) -> tuple[CellGridBounds, YX[slice]]:
212 subset_bounds = bounds[bbox]
213 start = subset_bounds.subgrid_start - bounds.subgrid_start
214 stop = subset_bounds.subgrid_stop - bounds.subgrid_start
215 return subset_bounds, YX(y=slice(start.i, stop.i), x=slice(start.j, stop.j))
218class CellPointSpreadFunctionSerializationModel(ArchiveTree):
219 """Model used to serialize CellPointSpreadFunction objects."""
221 SCHEMA_NAME: ClassVar[str] = "cell_psf"
222 SCHEMA_VERSION: ClassVar[str] = "1.0.0"
223 MIN_READ_VERSION: ClassVar[int] = 1
224 PUBLIC_TYPE: ClassVar[type] = CellPointSpreadFunction
226 array: ArrayReferenceModel | InlineArrayModel = pydantic.Field(
227 description=(
228 "A 4-d array of PSF kernel images with with shape "
229 "(n_cells_y, n_cells_x, psf_shape_y, psf_shape_x)."
230 )
231 )
232 bounds: CellGridBounds = pydantic.Field(
233 description=(
234 "Description of the cell grid and any missing cells. Array entries for "
235 "missing cells should be NaN."
236 )
237 )
239 def deserialize(
240 self, archive: InputArchive[Any], *, bbox: Box | None = None, **kwargs: Any
241 ) -> CellPointSpreadFunction:
242 if kwargs: 242 ↛ 243line 242 didn't jump to line 243 because the condition on line 242 was never true
243 raise InvalidParameterError(
244 f"Unrecognized parameters for CellPointSpreadFunction: {set(kwargs.keys())}."
245 )
246 bounds = self.bounds
247 if bbox is not None: 247 ↛ 248line 247 didn't jump to line 248 because the condition on line 247 was never true
248 bounds, slices = CellPointSpreadFunction._subset_impl(bounds, bbox)
249 array = archive.get_array(self.array, slices=slices)
250 else:
251 array = archive.get_array(self.array)
252 return CellPointSpreadFunction(array, bounds)