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-14 08:12 +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 

36 

37from astro_metadata_translator import ObservationInfo 

38from pydantic import BaseModel, Field 

39 

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 

54 

55 

56class ExtendedPsfCandidateInfo(BaseModel): 

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

58 

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

78 

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 

87 

88 def __str__(self) -> str: 

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

90 return f"ExtendedPsfCandidateInfo({attrs})" 

91 

92 __repr__ = __str__ 

93 

94 

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

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

97 

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

104 

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) 

117 

118 

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

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

121 

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

123 default_factory=list, 

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

125 ) 

126 

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 ) 

132 

133 

134class ExtendedPsfCandidate(MaskedImage): 

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

136 

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. 

157 

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

165 

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 ) 

188 

189 self._psf_kernel_image = psf_kernel_image 

190 self._star_info = star_info or ExtendedPsfCandidateInfo() 

191 

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 ) 

207 

208 def __str__(self) -> str: 

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

210 

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 ) 

216 

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 

223 

224 @property 

225 def star_info(self) -> ExtendedPsfCandidateInfo: 

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

227 return self._star_info 

228 

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 ) 

241 

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 ) 

257 

258 @staticmethod 

259 def _get_archive_tree_type[P: BaseModel]( 

260 pointer_type: type[P], 

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

262 return ExtendedPsfCandidateSerializationModel[pointer_type] 

263 

264 

265class ExtendedPsfCandidates(Sequence[ExtendedPsfCandidate]): 

266 """A collection of star cutouts. 

267 

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. 

274 

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

283 

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 } 

296 

297 def __len__(self): 

298 return len(self._candidates) 

299 

300 def __getitem__(self, index): 

301 if isinstance(index, slice): 

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

303 return self._candidates[index] 

304 

305 def __iter__(self): 

306 return iter(self._candidates) 

307 

308 def __str__(self) -> str: 

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

310 

311 __repr__ = __str__ 

312 

313 @property 

314 def metadata(self): 

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

316 return self._metadata 

317 

318 @property 

319 def ref_id_map(self): 

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

321 return self._ref_id_map 

322 

323 @classmethod 

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

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

326 

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 

334 

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

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

337 

338 Parameters 

339 ---------- 

340 filename 

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

342 """ 

343 fits.write(self, filename) 

344 

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 ) 

353 

354 @staticmethod 

355 def _get_archive_tree_type[P: BaseModel]( 

356 pointer_type: type[P], 

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

358 return ExtendedPsfCandidatesSerializationModel[pointer_type]