Coverage for python/lsst/pipe/tasks/extended_psf/extended_psf_candidates.py: 59%
104 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-03 08:13 +0000
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-03 08:13 +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, ClassVar
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 SCHEMA_NAME: ClassVar[str] = "extended_psf_candidate"
98 SCHEMA_VERSION: ClassVar[str] = "1.0.0"
99 MIN_READ_VERSION: ClassVar[int] = 1
101 psf_kernel_image: ImageSerializationModel[P] | None = Field(
102 default=None,
103 exclude_if=is_none,
104 description="Kernel image of the PSF at the cutout center.",
105 )
106 star_info: ExtendedPsfCandidateInfo = Field(
107 description="Information about the star in the cutout.",
108 )
110 def deserialize(self, archive: InputArchive[Any], *, bbox: Box | None = None) -> ExtendedPsfCandidate:
111 masked_image = super().deserialize(archive, bbox=bbox)
112 psf_kernel_image = (
113 self.psf_kernel_image.deserialize(archive) if self.psf_kernel_image is not None else None
114 )
115 return ExtendedPsfCandidate(
116 masked_image.image,
117 mask=masked_image.mask,
118 variance=masked_image.variance,
119 psf_kernel_image=psf_kernel_image,
120 star_info=self.star_info,
121 )._finish_deserialize(self)
124class ExtendedPsfCandidatesSerializationModel[P: BaseModel](ArchiveTree):
125 """A Pydantic model to represent serialized `ExtendedPsfCandidates`."""
127 SCHEMA_NAME: ClassVar[str] = "extended_psf_candidates"
128 SCHEMA_VERSION: ClassVar[str] = "1.0.0"
129 MIN_READ_VERSION: ClassVar[int] = 1
131 candidates: list[ExtendedPsfCandidateSerializationModel[P]] = Field(
132 default_factory=list,
133 description="The candidate cutouts in this collection.",
134 )
136 def deserialize(self, archive: InputArchive[Any]) -> ExtendedPsfCandidates:
137 return ExtendedPsfCandidates(
138 [candidate_model.deserialize(archive) for candidate_model in self.candidates],
139 metadata=self.metadata,
140 )
143class ExtendedPsfCandidate(MaskedImage):
144 """A cutout centered on a star, with associated metadata.
146 Parameters
147 ----------
148 image : `~lsst.images.Image`
149 The main data image for this star cutout.
150 mask : `~lsst.images.Mask`, optional
151 Bitmask that annotates the main image's pixels.
152 variance : `~lsst.images.Image`, optional
153 Per-pixel variance estimates for the image.
154 mask_schema : `~lsst.images.MaskSchema`, optional
155 Schema for the mask, required if a mask is provided.
156 projection : `~lsst.images.Projection`, optional
157 Projection to map pixels to the sky.
158 metadata : `dict` [`str`, `MetadataValue`], optional
159 Additional metadata to associate with this cutout.
160 psf_kernel_image : `~lsst.images.Image`, optional
161 Kernel image of the PSF at the cutout center.
162 star_info : `ExtendedPsfCandidateInfo`, optional
163 Information about the star in the cutout.
165 Attributes
166 ----------
167 psf_kernel_image : `~lsst.images.Image`
168 Kernel image of the PSF at the cutout center.
169 star_info : `ExtendedPsfCandidateInfo`
170 Information about the star in this cutout.
171 """
173 def __init__(
174 self,
175 image: Image,
176 *,
177 mask: Mask | None = None,
178 variance: Image | None = None,
179 mask_schema: MaskSchema | None = None,
180 projection: Projection | None = None,
181 metadata: dict[str, MetadataValue] | None = None,
182 psf_kernel_image: Image | None = None,
183 star_info: ExtendedPsfCandidateInfo | None = None,
184 ):
185 super().__init__(
186 image,
187 mask=mask,
188 variance=variance,
189 mask_schema=mask_schema,
190 projection=projection,
191 metadata=metadata,
192 )
194 self._psf_kernel_image = psf_kernel_image
195 self._star_info = star_info or ExtendedPsfCandidateInfo()
197 def __getitem__(self, bbox: Box | EllipsisType) -> ExtendedPsfCandidate:
198 if bbox is ...:
199 return self
200 super().__getitem__(bbox)
201 return self._transfer_metadata(
202 ExtendedPsfCandidate(
203 # Projection propagates from the image.
204 self.image[bbox],
205 mask=self.mask[bbox],
206 variance=self.variance[bbox],
207 psf_kernel_image=self.psf_kernel_image,
208 star_info=self.star_info,
209 ),
210 bbox=bbox,
211 )
213 def __str__(self) -> str:
214 return f"ExtendedPsfCandidate({self.image!s}, {list(self.mask.schema.names)}, {self.star_info})"
216 def __repr__(self) -> str:
217 return (
218 f"ExtendedPsfCandidate({self.image!r}, mask_schema={self.mask.schema!r}, "
219 f"star_info={self.star_info!r})"
220 )
222 @property
223 def psf_kernel_image(self) -> Image:
224 """Kernel image of the PSF at the cutout center."""
225 if self._psf_kernel_image is None:
226 raise RuntimeError("No PSF kernel image is attached to this ExtendedPsfCandidate.")
227 return self._psf_kernel_image
229 @property
230 def star_info(self) -> ExtendedPsfCandidateInfo:
231 """Return the ExtendedPsfCandidateInfo associated with this star."""
232 return self._star_info
234 def copy(self) -> ExtendedPsfCandidate:
235 """Deep-copy the star cutout, metadata, and star info."""
236 return self._transfer_metadata(
237 ExtendedPsfCandidate(
238 image=self._image.copy(),
239 mask=self._mask.copy(),
240 variance=self._variance.copy(),
241 psf_kernel_image=self._psf_kernel_image,
242 star_info=self._star_info.model_copy(),
243 ),
244 copy=True,
245 )
247 def serialize(self, archive: OutputArchive[Any]) -> ExtendedPsfCandidateSerializationModel:
248 masked_image_model = super().serialize(archive)
249 serialized_psf_kernel_image = (
250 archive.serialize_direct(
251 "psf_kernel_image",
252 functools.partial(self._psf_kernel_image.serialize, save_projection=False),
253 )
254 if self._psf_kernel_image is not None
255 else None
256 )
257 return ExtendedPsfCandidateSerializationModel(
258 **masked_image_model.model_dump(),
259 psf_kernel_image=serialized_psf_kernel_image,
260 star_info=self.star_info,
261 )
263 @staticmethod
264 def _get_archive_tree_type[P: BaseModel](
265 pointer_type: type[P],
266 ) -> type[ExtendedPsfCandidateSerializationModel[P]]:
267 return ExtendedPsfCandidateSerializationModel[pointer_type]
270class ExtendedPsfCandidates(Sequence[ExtendedPsfCandidate]):
271 """A collection of star cutouts.
273 Parameters
274 ----------
275 candidates : `Iterable` [`ExtendedPsfCandidate`]
276 Collection of `ExtendedPsfCandidate` instances.
277 metadata : `dict` [`str`, `MetadataValue`], optional
278 Global metadata associated with the collection.
280 Attributes
281 ----------
282 metadata : `dict` [`str`, `MetadataValue`]
283 Global metadata associated with the collection.
284 ref_id_map : `dict` [`int`, `ExtendedPsfCandidate`]
285 A mapping from reference IDs to `ExtendedPsfCandidate` objects.
286 Only includes candidates with valid reference IDs.
287 """
289 def __init__(
290 self,
291 candidates: Sequence[ExtendedPsfCandidate],
292 metadata: dict[str, MetadataValue] | None = None,
293 ):
294 self._candidates = list(candidates)
295 self._metadata = {} if metadata is None else dict(metadata)
296 self._ref_id_map = {
297 candidate.star_info.ref_id: candidate
298 for candidate in self
299 if candidate.star_info.ref_id is not None
300 }
302 def __len__(self):
303 return len(self._candidates)
305 def __getitem__(self, index):
306 if isinstance(index, slice):
307 return ExtendedPsfCandidates(self._candidates[index], metadata=self._metadata)
308 return self._candidates[index]
310 def __iter__(self):
311 return iter(self._candidates)
313 def __str__(self) -> str:
314 return f"ExtendedPsfCandidates(length={len(self)})"
316 __repr__ = __str__
318 @property
319 def metadata(self):
320 """Return the collection's global metadata as a dict."""
321 return self._metadata
323 @property
324 def ref_id_map(self):
325 """Map reference IDs to `ExtendedPsfCandidate` objects."""
326 return self._ref_id_map
328 @classmethod
329 def read_fits(cls, url: ResourcePathExpression) -> ExtendedPsfCandidates:
330 """Read a collection from a FITS file.
332 Parameters
333 ----------
334 url
335 URL of the file to read; may be any type supported by
336 `lsst.resources.ResourcePath`.
337 """
338 return fits.read(cls, url).deserialized
340 def write_fits(self, filename: str) -> None:
341 """Write the collection to a FITS file.
343 Parameters
344 ----------
345 filename
346 Name of the file to write to. Must not already exist.
347 """
348 fits.write(self, filename)
350 def serialize(self, archive: OutputArchive[Any]) -> ExtendedPsfCandidatesSerializationModel:
351 return ExtendedPsfCandidatesSerializationModel(
352 candidates=[
353 archive.serialize_direct(f"candidate_{index}", candidate.serialize)
354 for index, candidate in enumerate(self._candidates)
355 ],
356 metadata=self._metadata,
357 )
359 @staticmethod
360 def _get_archive_tree_type[P: BaseModel](
361 pointer_type: type[P],
362 ) -> type[ExtendedPsfCandidatesSerializationModel[P]]:
363 return ExtendedPsfCandidatesSerializationModel[pointer_type]