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