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