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-12 07:55 +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/>. 

21 

22from __future__ import annotations 

23 

24__all__ = ( 

25 "ExtendedPsfCandidateInfo", 

26 "ExtendedPsfCandidateSerializationModel", 

27 "ExtendedPsfCandidatesSerializationModel", 

28 "ExtendedPsfCandidate", 

29 "ExtendedPsfCandidates", 

30) 

31 

32import functools 

33from collections.abc import Sequence 

34from types import EllipsisType 

35from typing import Any, ClassVar 

36 

37from pydantic import BaseModel, Field 

38 

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 

53 

54 

55class ExtendedPsfCandidateInfo(BaseModel): 

56 """Information about a star in an `ExtendedPsfCandidate`. 

57 

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 """ 

77 

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 

86 

87 def __str__(self) -> str: 

88 attrs = ", ".join(f"{k}={v!r}" for k, v in self.__dict__.items()) 

89 return f"ExtendedPsfCandidateInfo({attrs})" 

90 

91 __repr__ = __str__ 

92 

93 

94class ExtendedPsfCandidate(MaskedImage): 

95 """A cutout centered on a star, with associated metadata. 

96 

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. 

115 

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 """ 

123 

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 ) 

144 

145 self._psf_kernel_image = psf_kernel_image 

146 self._star_info = star_info or ExtendedPsfCandidateInfo() 

147 

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 ) 

163 

164 def __str__(self) -> str: 

165 return f"ExtendedPsfCandidate({self.image!s}, {list(self.mask.schema.names)}, {self.star_info})" 

166 

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 ) 

172 

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 

179 

180 @property 

181 def star_info(self) -> ExtendedPsfCandidateInfo: 

182 """Return the ExtendedPsfCandidateInfo associated with this star.""" 

183 return self._star_info 

184 

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 ) 

197 

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 ) 

213 

214 @staticmethod 

215 def _get_archive_tree_type[P: BaseModel]( 

216 pointer_type: type[P], 

217 ) -> type[ExtendedPsfCandidateSerializationModel[P]]: 

218 return ExtendedPsfCandidateSerializationModel[pointer_type] 

219 

220 

221class ExtendedPsfCandidateSerializationModel[P: BaseModel](MaskedImageSerializationModel[P]): 

222 """A Pydantic model to represent a serialized `ExtendedPsfCandidate`.""" 

223 

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 

228 

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 ) 

237 

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) 

250 

251 

252class ExtendedPsfCandidates(Sequence[ExtendedPsfCandidate]): 

253 """A collection of star cutouts. 

254 

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. 

261 

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 """ 

270 

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 } 

283 

284 def __len__(self): 

285 return len(self._candidates) 

286 

287 def __getitem__(self, index): 

288 if isinstance(index, slice): 

289 return ExtendedPsfCandidates(self._candidates[index], metadata=self._metadata) 

290 return self._candidates[index] 

291 

292 def __iter__(self): 

293 return iter(self._candidates) 

294 

295 def __str__(self) -> str: 

296 return f"ExtendedPsfCandidates(length={len(self)})" 

297 

298 __repr__ = __str__ 

299 

300 @property 

301 def metadata(self): 

302 """Return the collection's global metadata as a dict.""" 

303 return self._metadata 

304 

305 @property 

306 def ref_id_map(self): 

307 """Map reference IDs to `ExtendedPsfCandidate` objects.""" 

308 return self._ref_id_map 

309 

310 @classmethod 

311 def read_fits(cls, url: ResourcePathExpression) -> ExtendedPsfCandidates: 

312 """Read a collection from a FITS file. 

313 

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 

321 

322 def write_fits(self, filename: str) -> None: 

323 """Write the collection to a FITS file. 

324 

325 Parameters 

326 ---------- 

327 filename 

328 Name of the file to write to. Must not already exist. 

329 """ 

330 fits.write(self, filename) 

331 

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 ) 

340 

341 @staticmethod 

342 def _get_archive_tree_type[P: BaseModel]( 

343 pointer_type: type[P], 

344 ) -> type[ExtendedPsfCandidatesSerializationModel[P]]: 

345 return ExtendedPsfCandidatesSerializationModel[pointer_type] 

346 

347 

348class ExtendedPsfCandidatesSerializationModel[P: BaseModel](ArchiveTree): 

349 """A Pydantic model to represent serialized `ExtendedPsfCandidates`.""" 

350 

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 

355 

356 candidates: list[ExtendedPsfCandidateSerializationModel[P]] = Field( 

357 default_factory=list, 

358 description="The candidate cutouts in this collection.", 

359 ) 

360 

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 )