Coverage for python/lsst/pipe/tasks/extended_psf/extended_psf_candidates.py: 57%
98 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-05-29 08:48 +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/>.
22from __future__ import annotations
24__all__ = (
25 "ExtendedPsfCandidateInfo",
26 "ExtendedPsfCandidateSerializationModel",
27 "ExtendedPsfCandidatesSerializationModel",
28 "ExtendedPsfCandidate",
29 "ExtendedPsfCandidates",
30)
32import functools
33from collections.abc import Sequence
34from types import EllipsisType
35from typing import Any
37from pydantic import BaseModel, Field
39from lsst.images import (
40 Box,
41 Image,
42 ImageSerializationModel,
43 Mask,
44 MaskedImage,
45 MaskedImageSerializationModel,
46 MaskSchema,
47 Projection,
48 fits,
49)
50from lsst.images.serialization import ArchiveTree, InputArchive, MetadataValue, OutputArchive, Quantity
51from lsst.images.utils import is_none
52from lsst.resources import ResourcePathExpression
55class ExtendedPsfCandidateInfo(BaseModel):
56 """Information about a star in an `ExtendedPsfCandidate`.
58 Attributes
59 ----------
60 visit : `int`, optional
61 The visit during which the star was observed.
62 detector : `int`, optional
63 The detector on which the star was observed.
64 ref_id : `int`, optional
65 The reference catalog ID for the star.
66 ref_mag : `float`, optional
67 The reference magnitude for the star.
68 position_x : `float`, optional
69 The x-coordinate of the star in the focal plane.
70 position_y : `float`, optional
71 The y-coordinate of the star in the focal plane.
72 focal_plane_radius : `~lsst.images.utils.Quantity`, optional
73 The radius of the star from the center of the focal plane.
74 focal_plane_angle : `~lsst.images.utils.Quantity`, optional
75 The angle of the star in the focal plane, measured from the +x axis.
76 """
78 visit: int | None = None
79 detector: int | None = None
80 ref_id: int | None = None
81 ref_mag: float | None = None
82 position_x: float | None = None
83 position_y: float | None = None
84 focal_plane_radius: Quantity | None = None
85 focal_plane_angle: Quantity | None = None
87 def __str__(self) -> str:
88 attrs = ", ".join(f"{k}={v!r}" for k, v in self.__dict__.items())
89 return f"ExtendedPsfCandidateInfo({attrs})"
91 __repr__ = __str__
94class ExtendedPsfCandidateSerializationModel[P: BaseModel](MaskedImageSerializationModel[P]):
95 """A Pydantic model to represent a serialized `ExtendedPsfCandidate`."""
97 psf_kernel_image: ImageSerializationModel[P] | None = Field(
98 default=None,
99 exclude_if=is_none,
100 description="Kernel image of the PSF at the cutout center.",
101 )
102 star_info: ExtendedPsfCandidateInfo = Field(
103 description="Information about the star in the cutout.",
104 )
106 def deserialize(self, archive: InputArchive[Any], *, bbox: Box | None = None) -> ExtendedPsfCandidate:
107 masked_image = super().deserialize(archive, bbox=bbox)
108 psf_kernel_image = (
109 self.psf_kernel_image.deserialize(archive) if self.psf_kernel_image is not None else None
110 )
111 return ExtendedPsfCandidate(
112 masked_image.image,
113 mask=masked_image.mask,
114 variance=masked_image.variance,
115 psf_kernel_image=psf_kernel_image,
116 star_info=self.star_info,
117 )._finish_deserialize(self)
120class ExtendedPsfCandidatesSerializationModel[P: BaseModel](ArchiveTree):
121 """A Pydantic model to represent serialized `ExtendedPsfCandidates`."""
123 candidates: list[ExtendedPsfCandidateSerializationModel[P]] = Field(
124 default_factory=list,
125 description="The candidate cutouts in this collection.",
126 )
128 def deserialize(self, archive: InputArchive[Any]) -> ExtendedPsfCandidates:
129 return ExtendedPsfCandidates(
130 [candidate_model.deserialize(archive) for candidate_model in self.candidates],
131 metadata=self.metadata,
132 )
135class ExtendedPsfCandidate(MaskedImage):
136 """A cutout centered on a star, with associated metadata.
138 Parameters
139 ----------
140 image : `~lsst.images.Image`
141 The main data image for this star cutout.
142 mask : `~lsst.images.Mask`, optional
143 Bitmask that annotates the main image's pixels.
144 variance : `~lsst.images.Image`, optional
145 Per-pixel variance estimates for the image.
146 mask_schema : `~lsst.images.MaskSchema`, optional
147 Schema for the mask, required if a mask is provided.
148 projection : `~lsst.images.Projection`, optional
149 Projection to map pixels to the sky.
150 metadata : `dict` [`str`, `MetadataValue`], optional
151 Additional metadata to associate with this cutout.
152 psf_kernel_image : `~lsst.images.Image`, optional
153 Kernel image of the PSF at the cutout center.
154 star_info : `ExtendedPsfCandidateInfo`, optional
155 Information about the star in the cutout.
157 Attributes
158 ----------
159 psf_kernel_image : `~lsst.images.Image`
160 Kernel image of the PSF at the cutout center.
161 star_info : `ExtendedPsfCandidateInfo`
162 Information about the star in this cutout.
163 """
165 def __init__(
166 self,
167 image: Image,
168 *,
169 mask: Mask | None = None,
170 variance: Image | None = None,
171 mask_schema: MaskSchema | None = None,
172 projection: Projection | None = None,
173 metadata: dict[str, MetadataValue] | None = None,
174 psf_kernel_image: Image | None = None,
175 star_info: ExtendedPsfCandidateInfo | None = None,
176 ):
177 super().__init__(
178 image,
179 mask=mask,
180 variance=variance,
181 mask_schema=mask_schema,
182 projection=projection,
183 metadata=metadata,
184 )
186 self._psf_kernel_image = psf_kernel_image
187 self._star_info = star_info or ExtendedPsfCandidateInfo()
189 def __getitem__(self, bbox: Box | EllipsisType) -> ExtendedPsfCandidate:
190 if bbox is ...:
191 return self
192 super().__getitem__(bbox)
193 return self._transfer_metadata(
194 ExtendedPsfCandidate(
195 # Projection propagates from the image.
196 self.image[bbox],
197 mask=self.mask[bbox],
198 variance=self.variance[bbox],
199 psf_kernel_image=self.psf_kernel_image,
200 star_info=self.star_info,
201 ),
202 bbox=bbox,
203 )
205 def __str__(self) -> str:
206 return f"ExtendedPsfCandidate({self.image!s}, {list(self.mask.schema.names)}, {self.star_info})"
208 def __repr__(self) -> str:
209 return (
210 f"ExtendedPsfCandidate({self.image!r}, mask_schema={self.mask.schema!r}, "
211 f"star_info={self.star_info!r})"
212 )
214 @property
215 def psf_kernel_image(self) -> Image:
216 """Kernel image of the PSF at the cutout center."""
217 if self._psf_kernel_image is None:
218 raise RuntimeError("No PSF kernel image is attached to this ExtendedPsfCandidate.")
219 return self._psf_kernel_image
221 @property
222 def star_info(self) -> ExtendedPsfCandidateInfo:
223 """Return the ExtendedPsfCandidateInfo associated with this star."""
224 return self._star_info
226 def copy(self) -> ExtendedPsfCandidate:
227 """Deep-copy the star cutout, metadata, and star info."""
228 return self._transfer_metadata(
229 ExtendedPsfCandidate(
230 image=self._image.copy(),
231 mask=self._mask.copy(),
232 variance=self._variance.copy(),
233 psf_kernel_image=self._psf_kernel_image,
234 star_info=self._star_info.model_copy(),
235 ),
236 copy=True,
237 )
239 def serialize(self, archive: OutputArchive[Any]) -> ExtendedPsfCandidateSerializationModel:
240 masked_image_model = super().serialize(archive)
241 serialized_psf_kernel_image = (
242 archive.serialize_direct(
243 "psf_kernel_image",
244 functools.partial(self._psf_kernel_image.serialize, save_projection=False),
245 )
246 if self._psf_kernel_image is not None
247 else None
248 )
249 return ExtendedPsfCandidateSerializationModel(
250 **masked_image_model.model_dump(),
251 psf_kernel_image=serialized_psf_kernel_image,
252 star_info=self.star_info,
253 )
255 @staticmethod
256 def _get_archive_tree_type[P: BaseModel](
257 pointer_type: type[P],
258 ) -> type[ExtendedPsfCandidateSerializationModel[P]]:
259 return ExtendedPsfCandidateSerializationModel[pointer_type]
262class ExtendedPsfCandidates(Sequence[ExtendedPsfCandidate]):
263 """A collection of star cutouts.
265 Parameters
266 ----------
267 candidates : `Iterable` [`ExtendedPsfCandidate`]
268 Collection of `ExtendedPsfCandidate` instances.
269 metadata : `dict` [`str`, `MetadataValue`], optional
270 Global metadata associated with the collection.
272 Attributes
273 ----------
274 metadata : `dict` [`str`, `MetadataValue`]
275 Global metadata associated with the collection.
276 ref_id_map : `dict` [`int`, `ExtendedPsfCandidate`]
277 A mapping from reference IDs to `ExtendedPsfCandidate` objects.
278 Only includes candidates with valid reference IDs.
279 """
281 def __init__(
282 self,
283 candidates: Sequence[ExtendedPsfCandidate],
284 metadata: dict[str, MetadataValue] | None = None,
285 ):
286 self._candidates = list(candidates)
287 self._metadata = {} if metadata is None else dict(metadata)
288 self._ref_id_map = {
289 candidate.star_info.ref_id: candidate
290 for candidate in self
291 if candidate.star_info.ref_id is not None
292 }
294 def __len__(self):
295 return len(self._candidates)
297 def __getitem__(self, index):
298 if isinstance(index, slice):
299 return ExtendedPsfCandidates(self._candidates[index], metadata=self._metadata)
300 return self._candidates[index]
302 def __iter__(self):
303 return iter(self._candidates)
305 def __str__(self) -> str:
306 return f"ExtendedPsfCandidates(length={len(self)})"
308 __repr__ = __str__
310 @property
311 def metadata(self):
312 """Return the collection's global metadata as a dict."""
313 return self._metadata
315 @property
316 def ref_id_map(self):
317 """Map reference IDs to `ExtendedPsfCandidate` objects."""
318 return self._ref_id_map
320 @classmethod
321 def read_fits(cls, url: ResourcePathExpression) -> ExtendedPsfCandidates:
322 """Read a collection from a FITS file.
324 Parameters
325 ----------
326 url
327 URL of the file to read; may be any type supported by
328 `lsst.resources.ResourcePath`.
329 """
330 return fits.read(cls, url).deserialized
332 def write_fits(self, filename: str) -> None:
333 """Write the collection to a FITS file.
335 Parameters
336 ----------
337 filename
338 Name of the file to write to. Must not already exist.
339 """
340 fits.write(self, filename)
342 def serialize(self, archive: OutputArchive[Any]) -> ExtendedPsfCandidatesSerializationModel:
343 return ExtendedPsfCandidatesSerializationModel(
344 candidates=[
345 archive.serialize_direct(f"candidate_{index}", candidate.serialize)
346 for index, candidate in enumerate(self._candidates)
347 ],
348 metadata=self._metadata,
349 )
351 @staticmethod
352 def _get_archive_tree_type[P: BaseModel](
353 pointer_type: type[P],
354 ) -> type[ExtendedPsfCandidatesSerializationModel[P]]:
355 return ExtendedPsfCandidatesSerializationModel[pointer_type]