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

101 statements  

« prev     ^ index     » next       coverage.py v7.13.5, created at 2026-04-28 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/>. 

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 

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

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

108 

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

110 default_factory=list, 

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

112 ) 

113 

114 

115class ExtendedPsfCandidate(MaskedImage): 

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

117 

118 Parameters 

119 ---------- 

120 image : `~lsst.images.Image` 

121 The main data image for this star cutout. 

122 mask : `~lsst.images.Mask`, optional 

123 Bitmask that annotates the main image's pixels. 

124 variance : `~lsst.images.Image`, optional 

125 Per-pixel variance estimates for the image. 

126 mask_schema : `~lsst.images.MaskSchema`, optional 

127 Schema for the mask, required if a mask is provided. 

128 projection : `~lsst.images.Projection`, optional 

129 Projection to map pixels to the sky. 

130 obs_info : `~astro_metadata_translator.ObservationInfo`, optional 

131 Standardized description of visit metadata. 

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

133 Additional metadata to associate with this cutout. 

134 psf_kernel_image : `~lsst.images.Image`, optional 

135 Kernel image of the PSF at the cutout center. 

136 star_info : `ExtendedPsfCandidateInfo`, optional 

137 Information about the star in the cutout. 

138 

139 Attributes 

140 ---------- 

141 psf_kernel_image : `~lsst.images.Image` 

142 Kernel image of the PSF at the cutout center. 

143 star_info : `ExtendedPsfCandidateInfo` 

144 Information about the star in this cutout. 

145 """ 

146 

147 def __init__( 

148 self, 

149 image: Image, 

150 *, 

151 mask: Mask | None = None, 

152 variance: Image | None = None, 

153 mask_schema: MaskSchema | None = None, 

154 projection: Projection | None = None, 

155 obs_info: ObservationInfo | None = None, 

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

157 psf_kernel_image: Image | None = None, 

158 star_info: ExtendedPsfCandidateInfo | None = None, 

159 ): 

160 super().__init__( 

161 image, 

162 mask=mask, 

163 variance=variance, 

164 mask_schema=mask_schema, 

165 projection=projection, 

166 obs_info=obs_info, 

167 metadata=metadata, 

168 ) 

169 

170 self._psf_kernel_image = psf_kernel_image 

171 self._star_info = star_info or ExtendedPsfCandidateInfo() 

172 

173 def __getitem__(self, bbox: Box | EllipsisType) -> ExtendedPsfCandidate: 

174 if bbox is ...: 

175 return self 

176 super().__getitem__(bbox) 

177 return self._transfer_metadata( 

178 ExtendedPsfCandidate( 

179 # Projection and obs_info propagate from the image. 

180 self.image[bbox], 

181 mask=self.mask[bbox], 

182 variance=self.variance[bbox], 

183 psf_kernel_image=self.psf_kernel_image, 

184 star_info=self.star_info, 

185 ), 

186 bbox=bbox, 

187 ) 

188 

189 def __str__(self) -> str: 

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

191 

192 def __repr__(self) -> str: 

193 return ( 

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

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

196 ) 

197 

198 @property 

199 def psf_kernel_image(self) -> Image: 

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

201 if self._psf_kernel_image is None: 

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

203 return self._psf_kernel_image 

204 

205 @property 

206 def star_info(self) -> ExtendedPsfCandidateInfo: 

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

208 return self._star_info 

209 

210 def copy(self) -> ExtendedPsfCandidate: 

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

212 return self._transfer_metadata( 

213 ExtendedPsfCandidate( 

214 image=self._image.copy(), 

215 mask=self._mask.copy(), 

216 variance=self._variance.copy(), 

217 psf_kernel_image=self._psf_kernel_image, 

218 star_info=self._star_info.model_copy(), 

219 ), 

220 copy=True, 

221 ) 

222 

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

224 masked_image_model = super().serialize(archive) 

225 serialized_psf_kernel_image = ( 

226 archive.serialize_direct( 

227 "psf_kernel_image", 

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

229 ) 

230 if self._psf_kernel_image is not None 

231 else None 

232 ) 

233 return ExtendedPsfCandidateSerializationModel( 

234 **masked_image_model.model_dump(), 

235 psf_kernel_image=serialized_psf_kernel_image, 

236 star_info=self.star_info, 

237 ) 

238 

239 @staticmethod 

240 def _get_archive_tree_type[P: BaseModel]( 

241 pointer_type: type[P], 

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

243 return ExtendedPsfCandidateSerializationModel[pointer_type] 

244 

245 @staticmethod 

246 def deserialize( 

247 model: ExtendedPsfCandidateSerializationModel[Any], 

248 archive: InputArchive[Any], 

249 *, 

250 bbox: Box | None = None, 

251 ) -> ExtendedPsfCandidate: 

252 masked_image = MaskedImage.deserialize(model, archive, bbox=bbox) 

253 psf_kernel_image = ( 

254 Image.deserialize(model.psf_kernel_image, archive) if model.psf_kernel_image is not None else None 

255 ) 

256 return ExtendedPsfCandidate( 

257 masked_image.image, 

258 mask=masked_image.mask, 

259 variance=masked_image.variance, 

260 psf_kernel_image=psf_kernel_image, 

261 star_info=model.star_info, 

262 )._finish_deserialize(model) 

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 deserialize( 

356 model: ExtendedPsfCandidatesSerializationModel[Any], 

357 archive: InputArchive[Any], 

358 ) -> ExtendedPsfCandidates: 

359 return ExtendedPsfCandidates( 

360 [ 

361 ExtendedPsfCandidate.deserialize(candidate_model, archive) 

362 for candidate_model in model.candidates 

363 ], 

364 metadata=model.metadata, 

365 ) 

366 

367 @staticmethod 

368 def _get_archive_tree_type[P: BaseModel]( 

369 pointer_type: type[P], 

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

371 return ExtendedPsfCandidatesSerializationModel[pointer_type]