Coverage for python/lsst/pipe/tasks/extended_psf/extended_psf_candidates.py: 90%

104 statements  

« prev     ^ index     » next       coverage.py v7.15.2, created at 2026-07-23 09:30 +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 bbox, _ = self._handle_getitem_args(bbox) 

150 return self._transfer_metadata( 

151 ExtendedPsfCandidate( 

152 # Projection propagates from the image. 

153 self.image[bbox], 

154 mask=self.mask[bbox], 

155 variance=self.variance[bbox], 

156 psf_kernel_image=self.psf_kernel_image, 

157 star_info=self.star_info, 

158 ), 

159 bbox=bbox, 

160 ) 

161 

162 def __str__(self) -> str: 

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

164 

165 def __repr__(self) -> str: 

166 return ( 

167 f"ExtendedPsfCandidate({self.image!r}, mask_schema={self.mask.schema!r}, " 

168 f"star_info={self.star_info!r})" 

169 ) 

170 

171 @property 

172 def psf_kernel_image(self) -> Image: 

173 """Kernel image of the PSF at the cutout center.""" 

174 if self._psf_kernel_image is None: 174 ↛ 175line 174 didn't jump to line 175 because the condition on line 174 was never true

175 raise RuntimeError("No PSF kernel image is attached to this ExtendedPsfCandidate.") 

176 return self._psf_kernel_image 

177 

178 @property 

179 def star_info(self) -> ExtendedPsfCandidateInfo: 

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

181 return self._star_info 

182 

183 def copy(self) -> ExtendedPsfCandidate: 

184 """Deep-copy the star cutout, metadata, and star info.""" 

185 return self._transfer_metadata( 

186 ExtendedPsfCandidate( 

187 image=self._image.copy(), 

188 mask=self._mask.copy(), 

189 variance=self._variance.copy(), 

190 psf_kernel_image=self._psf_kernel_image, 

191 star_info=self._star_info.model_copy(), 

192 ), 

193 copy=True, 

194 ) 

195 

196 def serialize(self, archive: OutputArchive[Any]) -> ExtendedPsfCandidateSerializationModel: 

197 masked_image_model = super().serialize(archive) 

198 serialized_psf_kernel_image = ( 

199 archive.serialize_direct( 

200 "psf_kernel_image", 

201 functools.partial(self._psf_kernel_image.serialize, save_projection=False), 

202 ) 

203 if self._psf_kernel_image is not None 

204 else None 

205 ) 

206 return ExtendedPsfCandidateSerializationModel( 

207 **masked_image_model.model_dump(), 

208 psf_kernel_image=serialized_psf_kernel_image, 

209 star_info=self.star_info, 

210 ) 

211 

212 @staticmethod 

213 def _get_archive_tree_type[P: BaseModel]( 

214 pointer_type: type[P], 

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

216 return ExtendedPsfCandidateSerializationModel[pointer_type] 

217 

218 

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

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

221 

222 SCHEMA_NAME: ClassVar[str] = "extended_psf_candidate" 

223 SCHEMA_VERSION: ClassVar[str] = "1.0.0" 

224 MIN_READ_VERSION: ClassVar[int] = 1 

225 PUBLIC_TYPE: ClassVar[type] = ExtendedPsfCandidate 

226 

227 psf_kernel_image: ImageSerializationModel[P] | None = Field( 

228 default=None, 

229 exclude_if=is_none, 

230 description="Kernel image of the PSF at the cutout center.", 

231 ) 

232 star_info: ExtendedPsfCandidateInfo = Field( 

233 description="Information about the star in the cutout.", 

234 ) 

235 

236 def deserialize(self, archive: InputArchive[Any], *, bbox: Box | None = None) -> ExtendedPsfCandidate: 

237 masked_image = super().deserialize(archive, bbox=bbox) 

238 psf_kernel_image = ( 

239 self.psf_kernel_image.deserialize(archive) if self.psf_kernel_image is not None else None 

240 ) 

241 return ExtendedPsfCandidate( 

242 masked_image.image, 

243 mask=masked_image.mask, 

244 variance=masked_image.variance, 

245 psf_kernel_image=psf_kernel_image, 

246 star_info=self.star_info, 

247 )._finish_deserialize(self) 

248 

249 

250class ExtendedPsfCandidates(Sequence[ExtendedPsfCandidate]): 

251 """A collection of star cutouts. 

252 

253 Parameters 

254 ---------- 

255 candidates : `Iterable` [`ExtendedPsfCandidate`] 

256 Collection of `ExtendedPsfCandidate` instances. 

257 metadata : `dict` [`str`, `MetadataValue`], optional 

258 Global metadata associated with the collection. 

259 

260 Attributes 

261 ---------- 

262 metadata : `dict` [`str`, `MetadataValue`] 

263 Global metadata associated with the collection. 

264 ref_id_map : `dict` [`int`, `ExtendedPsfCandidate`] 

265 A mapping from reference IDs to `ExtendedPsfCandidate` objects. 

266 Only includes candidates with valid reference IDs. 

267 """ 

268 

269 def __init__( 

270 self, 

271 candidates: Sequence[ExtendedPsfCandidate], 

272 metadata: dict[str, MetadataValue] | None = None, 

273 ): 

274 self._candidates = list(candidates) 

275 self._metadata = {} if metadata is None else dict(metadata) 

276 self._ref_id_map = { 

277 candidate.star_info.ref_id: candidate 

278 for candidate in self 

279 if candidate.star_info.ref_id is not None 

280 } 

281 

282 def __len__(self): 

283 return len(self._candidates) 

284 

285 def __getitem__(self, index): 

286 if isinstance(index, slice): 

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

288 return self._candidates[index] 

289 

290 def __iter__(self): 

291 return iter(self._candidates) 

292 

293 def __str__(self) -> str: 

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

295 

296 __repr__ = __str__ 

297 

298 @property 

299 def metadata(self): 

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

301 return self._metadata 

302 

303 @property 

304 def ref_id_map(self): 

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

306 return self._ref_id_map 

307 

308 @classmethod 

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

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

311 

312 Parameters 

313 ---------- 

314 url 

315 URL of the file to read; may be any type supported by 

316 `lsst.resources.ResourcePath`. 

317 """ 

318 return fits.read(cls, url).deserialized 

319 

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

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

322 

323 Parameters 

324 ---------- 

325 filename 

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

327 """ 

328 fits.write(self, filename) 

329 

330 def serialize(self, archive: OutputArchive[Any]) -> ExtendedPsfCandidatesSerializationModel: 

331 return ExtendedPsfCandidatesSerializationModel( 

332 candidates=[ 

333 archive.serialize_direct(f"candidate_{index}", candidate.serialize) 

334 for index, candidate in enumerate(self._candidates) 

335 ], 

336 metadata=self._metadata, 

337 ) 

338 

339 @staticmethod 

340 def _get_archive_tree_type[P: BaseModel]( 

341 pointer_type: type[P], 

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

343 return ExtendedPsfCandidatesSerializationModel[pointer_type] 

344 

345 

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

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

348 

349 SCHEMA_NAME: ClassVar[str] = "extended_psf_candidates" 

350 SCHEMA_VERSION: ClassVar[str] = "1.0.0" 

351 MIN_READ_VERSION: ClassVar[int] = 1 

352 PUBLIC_TYPE: ClassVar[type] = ExtendedPsfCandidates 

353 

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

355 default_factory=list, 

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

357 ) 

358 

359 def deserialize(self, archive: InputArchive[Any]) -> ExtendedPsfCandidates: 

360 return ExtendedPsfCandidates( 

361 [candidate_model.deserialize(archive) for candidate_model in self.candidates], 

362 metadata=self.metadata, 

363 )