Coverage for python/lsst/meas/algorithms/stamps.py: 24%

169 statements  

« prev     ^ index     » next       coverage.py v7.2.5, created at 2023-05-23 03:35 -0700

1# This file is part of meas_algorithms. 

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 

22"""Collection of small images (stamps).""" 

23 

24__all__ = ["Stamp", "Stamps", "StampsBase", "writeFits", "readFitsWithOptions"] 

25 

26import abc 

27from collections.abc import Sequence 

28from dataclasses import dataclass 

29 

30import numpy as np 

31from lsst.afw.fits import Fits, readMetadata 

32from lsst.afw.image import ImageFitsReader, MaskedImageF, MaskFitsReader 

33from lsst.afw.table.io import InputArchive, OutputArchive, Persistable 

34from lsst.daf.base import PropertyList 

35from lsst.geom import Angle, Box2I, Extent2I, Point2I, SpherePoint, degrees 

36from lsst.utils import doImport 

37from lsst.utils.introspection import get_full_type_name 

38 

39 

40def writeFits(filename, stamps, metadata, type_name, write_mask, write_variance, write_archive=False): 

41 """Write a single FITS file containing all stamps. 

42 

43 Parameters 

44 ---------- 

45 filename : `str` 

46 A string indicating the output filename 

47 stamps : iterable of `BaseStamp` 

48 An iterable of Stamp objects 

49 metadata : `PropertyList` 

50 A collection of key, value metadata pairs to be 

51 written to the primary header 

52 type_name : `str` 

53 Python type name of the StampsBase subclass to use 

54 write_mask : `bool` 

55 Write the mask data to the output file? 

56 write_variance : `bool` 

57 Write the variance data to the output file? 

58 write_archive : `bool`, optional 

59 Write an archive to store Persistables along with each stamp? 

60 Default: ``False``. 

61 """ 

62 metadata["HAS_MASK"] = write_mask 

63 metadata["HAS_VARIANCE"] = write_variance 

64 metadata["HAS_ARCHIVE"] = write_archive 

65 metadata["N_STAMPS"] = len(stamps) 

66 metadata["STAMPCLS"] = type_name 

67 # Record version number in case of future code changes 

68 metadata["VERSION"] = 1 

69 # create primary HDU with global metadata 

70 fitsFile = Fits(filename, "w") 

71 fitsFile.createEmpty() 

72 # Store Persistables in an OutputArchive and write it 

73 if write_archive: 

74 oa = OutputArchive() 

75 archive_ids = [oa.put(stamp.archive_element) for stamp in stamps] 

76 metadata["ARCHIVE_IDS"] = archive_ids 

77 fitsFile.writeMetadata(metadata) 

78 oa.writeFits(fitsFile) 

79 else: 

80 fitsFile.writeMetadata(metadata) 

81 fitsFile.closeFile() 

82 # add all pixel data optionally writing mask and variance information 

83 for i, stamp in enumerate(stamps): 

84 metadata = PropertyList() 

85 # EXTVER should be 1-based, the index from enumerate is 0-based 

86 metadata.update({"EXTVER": i + 1, "EXTNAME": "IMAGE"}) 

87 stamp.stamp_im.getImage().writeFits(filename, metadata=metadata, mode="a") 

88 if write_mask: 

89 metadata = PropertyList() 

90 metadata.update({"EXTVER": i + 1, "EXTNAME": "MASK"}) 

91 stamp.stamp_im.getMask().writeFits(filename, metadata=metadata, mode="a") 

92 if write_variance: 

93 metadata = PropertyList() 

94 metadata.update({"EXTVER": i + 1, "EXTNAME": "VARIANCE"}) 

95 stamp.stamp_im.getVariance().writeFits(filename, metadata=metadata, mode="a") 

96 return None 

97 

98 

99def readFitsWithOptions(filename, stamp_factory, options): 

100 """Read stamps from FITS file, allowing for only a subregion of the stamps 

101 to be read. 

102 

103 Parameters 

104 ---------- 

105 filename : `str` 

106 A string indicating the file to read 

107 stamp_factory : classmethod 

108 A factory function defined on a dataclass for constructing 

109 stamp objects a la `~lsst.meas.alrogithm.Stamp` 

110 options : `PropertyList` or `dict` 

111 A collection of parameters. If it contains a bounding box 

112 (``bbox`` key), or if certain other keys (``llcX``, ``llcY``, 

113 ``width``, ``height``) are available for one to be constructed, 

114 the bounding box is passed to the ``FitsReader`` in order to 

115 return a sub-image. 

116 

117 Returns 

118 ------- 

119 stamps : `list` of dataclass objects like `Stamp`, PropertyList 

120 A tuple of a list of `Stamp`-like objects 

121 metadata : `PropertyList` 

122 The metadata 

123 """ 

124 # extract necessary info from metadata 

125 metadata = readMetadata(filename, hdu=0) 

126 nStamps = metadata["N_STAMPS"] 

127 has_archive = metadata["HAS_ARCHIVE"] 

128 if has_archive: 

129 archive_ids = metadata.getArray("ARCHIVE_IDS") 

130 with Fits(filename, "r") as f: 

131 nExtensions = f.countHdus() 

132 # check if a bbox was provided 

133 kwargs = {} 

134 if options: 

135 # gen3 API 

136 if "bbox" in options.keys(): 

137 kwargs["bbox"] = options["bbox"] 

138 # gen2 API 

139 elif "llcX" in options.keys(): 

140 llcX = options["llcX"] 

141 llcY = options["llcY"] 

142 width = options["width"] 

143 height = options["height"] 

144 bbox = Box2I(Point2I(llcX, llcY), Extent2I(width, height)) 

145 kwargs["bbox"] = bbox 

146 stamp_parts = {} 

147 # We need to be careful because nExtensions includes the primary HDU. 

148 for idx in range(nExtensions - 1): 

149 md = readMetadata(filename, hdu=idx + 1) 

150 if md["EXTNAME"] in ("IMAGE", "VARIANCE"): 

151 reader = ImageFitsReader(filename, hdu=idx + 1) 

152 elif md["EXTNAME"] == "MASK": 

153 reader = MaskFitsReader(filename, hdu=idx + 1) 

154 elif md["EXTNAME"] == "ARCHIVE_INDEX": 

155 f.setHdu(idx + 1) 

156 archive = InputArchive.readFits(f) 

157 continue 

158 elif md["EXTTYPE"] == "ARCHIVE_DATA": 

159 continue 

160 else: 

161 raise ValueError(f"Unknown extension type: {md['EXTNAME']}") 

162 stamp_parts.setdefault(md["EXTVER"], {})[md["EXTNAME"].lower()] = reader.read(**kwargs) 

163 if len(stamp_parts) != nStamps: 

164 raise ValueError( 

165 f"Number of stamps read ({len(stamp_parts)}) does not agree with the " 

166 f"number of stamps recorded in the metadata ({nStamps})." 

167 ) 

168 # construct stamps themselves 

169 stamps = [] 

170 for k in range(nStamps): 

171 # Need to increment by one since EXTVER starts at 1 

172 maskedImage = MaskedImageF(**stamp_parts[k + 1]) 

173 archive_element = archive.get(archive_ids[k]) if has_archive else None 

174 stamps.append(stamp_factory(maskedImage, metadata, k, archive_element)) 

175 

176 return stamps, metadata 

177 

178 

179@dataclass 

180class AbstractStamp(abc.ABC): 

181 """Single abstract stamp. 

182 

183 Parameters 

184 ---------- 

185 Inherit from this class to add metadata to the stamp. 

186 """ 

187 

188 @classmethod 

189 @abc.abstractmethod 

190 def factory(cls, stamp_im, metadata, index, archive_element=None): 

191 """This method is needed to service the FITS reader. We need a standard 

192 interface to construct objects like this. Parameters needed to 

193 construct this object are passed in via a metadata dictionary and then 

194 passed to the constructor of this class. 

195 

196 Parameters 

197 ---------- 

198 stamp : `~lsst.afw.image.MaskedImage` 

199 Pixel data to pass to the constructor 

200 metadata : `dict` 

201 Dictionary containing the information 

202 needed by the constructor. 

203 idx : `int` 

204 Index into the lists in ``metadata`` 

205 archive_element : `~lsst.afw.table.io.Persistable`, optional 

206 Archive element (e.g. Transform or WCS) associated with this stamp. 

207 

208 Returns 

209 ------- 

210 stamp : `AbstractStamp` 

211 An instance of this class 

212 """ 

213 raise NotImplementedError 

214 

215 

216@dataclass 

217class Stamp(AbstractStamp): 

218 """Single stamp. 

219 

220 Parameters 

221 ---------- 

222 stamp_im : `~lsst.afw.image.MaskedImageF` 

223 The actual pixel values for the postage stamp. 

224 archive_element : `~lsst.afw.table.io.Persistable` or `None`, optional 

225 Archive element (e.g. Transform or WCS) associated with this stamp. 

226 position : `~lsst.geom.SpherePoint` or `None`, optional 

227 Position of the center of the stamp. Note the user must keep track of 

228 the coordinate system. 

229 """ 

230 

231 stamp_im: MaskedImageF 

232 archive_element: Persistable | None = None 

233 position: SpherePoint | None = SpherePoint(Angle(np.nan), Angle(np.nan)) 

234 

235 @classmethod 

236 def factory(cls, stamp_im, metadata, index, archive_element=None): 

237 """This method is needed to service the FITS reader. We need a standard 

238 interface to construct objects like this. Parameters needed to 

239 construct this object are passed in via a metadata dictionary and then 

240 passed to the constructor of this class. If lists of values are passed 

241 with the following keys, they will be passed to the constructor, 

242 otherwise dummy values will be passed: RA_DEG, DEC_DEG. They should 

243 each point to lists of values. 

244 

245 Parameters 

246 ---------- 

247 stamp : `~lsst.afw.image.MaskedImage` 

248 Pixel data to pass to the constructor 

249 metadata : `dict` 

250 Dictionary containing the information 

251 needed by the constructor. 

252 idx : `int` 

253 Index into the lists in ``metadata`` 

254 archive_element : `~lsst.afw.table.io.Persistable`, optional 

255 Archive element (e.g. Transform or WCS) associated with this stamp. 

256 

257 Returns 

258 ------- 

259 stamp : `Stamp` 

260 An instance of this class 

261 """ 

262 if "RA_DEG" in metadata and "DEC_DEG" in metadata: 

263 return cls( 

264 stamp_im=stamp_im, 

265 archive_element=archive_element, 

266 position=SpherePoint( 

267 Angle(metadata.getArray("RA_DEG")[index], degrees), 

268 Angle(metadata.getArray("DEC_DEG")[index], degrees), 

269 ), 

270 ) 

271 else: 

272 return cls( 

273 stamp_im=stamp_im, 

274 archive_element=archive_element, 

275 position=SpherePoint(Angle(np.nan), Angle(np.nan)), 

276 ) 

277 

278 

279class StampsBase(abc.ABC, Sequence): 

280 """Collection of stamps and associated metadata. 

281 

282 Parameters 

283 ---------- 

284 stamps : iterable 

285 This should be an iterable of dataclass objects 

286 a la ``~lsst.meas.algorithms.Stamp``. 

287 metadata : `~lsst.daf.base.PropertyList`, optional 

288 Metadata associated with the objects within the stamps. 

289 use_mask : `bool`, optional 

290 If ``True`` read and write the mask data. Default ``True``. 

291 use_variance : `bool`, optional 

292 If ``True`` read and write the variance data. Default ``True``. 

293 use_archive : `bool`, optional 

294 If ``True``, read and write an Archive that contains a Persistable 

295 associated with each stamp, for example a Transform or a WCS. 

296 Default ``False``. 

297 

298 Notes 

299 ----- 

300 A butler can be used to read only a part of the stamps, 

301 specified by a bbox: 

302 

303 >>> starSubregions = butler.get( 

304 "brightStarStamps", 

305 dataId, 

306 parameters={"bbox": bbox} 

307 ) 

308 """ 

309 

310 def __init__(self, stamps, metadata=None, use_mask=True, use_variance=True, use_archive=False): 

311 for stamp in stamps: 

312 if not isinstance(stamp, AbstractStamp): 

313 raise ValueError(f"The entries in stamps must inherit from AbstractStamp. Got {type(stamp)}.") 

314 self._stamps = stamps 

315 self._metadata = PropertyList() if metadata is None else metadata.deepCopy() 

316 self.use_mask = use_mask 

317 self.use_variance = use_variance 

318 self.use_archive = use_archive 

319 

320 @classmethod 

321 def readFits(cls, filename): 

322 """Build an instance of this class from a file. 

323 

324 Parameters 

325 ---------- 

326 filename : `str` 

327 Name of the file to read 

328 """ 

329 

330 return cls.readFitsWithOptions(filename, None) 

331 

332 @classmethod 

333 def readFitsWithOptions(cls, filename, options): 

334 """Build an instance of this class with options. 

335 

336 Parameters 

337 ---------- 

338 filename : `str` 

339 Name of the file to read 

340 options : `PropertyList` 

341 Collection of metadata parameters 

342 """ 

343 # To avoid problems since this is no longer an abstract method. 

344 # TO-DO: Consider refactoring this method. This class check was added 

345 # to allow the butler formatter to use a generic type but still end up 

346 # giving the correct type back, ensuring that the abstract base class 

347 # is not used by mistake. Perhaps this logic can be optimised. 

348 if cls is not StampsBase: 

349 raise NotImplementedError(f"Please implement specific FITS reader for class {cls}") 

350 

351 # Load metadata to get class 

352 metadata = readMetadata(filename, hdu=0) 

353 type_name = metadata.get("STAMPCLS") 

354 if type_name is None: 

355 raise RuntimeError( 

356 f"No class name in file {filename}. Unable to instantiate correct stamps subclass. " 

357 "Is this an old version format Stamps file?" 

358 ) 

359 

360 # Import class and override `cls` 

361 stamp_type = doImport(type_name) 

362 cls = stamp_type 

363 

364 return cls.readFitsWithOptions(filename, options) 

365 

366 @abc.abstractmethod 

367 def _refresh_metadata(self): 

368 """Make sure metadata is up to date, as this object can be extended.""" 

369 raise NotImplementedError 

370 

371 def writeFits(self, filename): 

372 """Write this object to a file. 

373 

374 Parameters 

375 ---------- 

376 filename : `str` 

377 Name of file to write. 

378 """ 

379 self._refresh_metadata() 

380 type_name = get_full_type_name(self) 

381 writeFits( 

382 filename, 

383 self._stamps, 

384 self._metadata, 

385 type_name, 

386 self.use_mask, 

387 self.use_variance, 

388 self.use_archive, 

389 ) 

390 

391 def __len__(self): 

392 return len(self._stamps) 

393 

394 def __getitem__(self, index): 

395 return self._stamps[index] 

396 

397 def __iter__(self): 

398 return iter(self._stamps) 

399 

400 def getMaskedImages(self): 

401 """Retrieve star images. 

402 

403 Returns 

404 ------- 

405 maskedImages : 

406 `list` [`~lsst.afw.image.MaskedImageF`] 

407 """ 

408 return [stamp.stamp_im for stamp in self._stamps] 

409 

410 def getArchiveElements(self): 

411 """Retrieve archive elements associated with each stamp. 

412 

413 Returns 

414 ------- 

415 archiveElements : 

416 `list` [`~lsst.afw.table.io.Persistable`] 

417 """ 

418 return [stamp.archive_element for stamp in self._stamps] 

419 

420 @property 

421 def metadata(self): 

422 return self._metadata 

423 

424 

425class Stamps(StampsBase): 

426 def _refresh_metadata(self): 

427 positions = self.getPositions() 

428 self._metadata["RA_DEG"] = [p.getRa().asDegrees() for p in positions] 

429 self._metadata["DEC_DEG"] = [p.getDec().asDegrees() for p in positions] 

430 

431 def getPositions(self): 

432 return [s.position for s in self._stamps] 

433 

434 def append(self, item): 

435 """Add an additional stamp. 

436 

437 Parameters 

438 ---------- 

439 item : `Stamp` 

440 Stamp object to append. 

441 """ 

442 if not isinstance(item, Stamp): 

443 raise ValueError("Objects added must be a Stamp object.") 

444 self._stamps.append(item) 

445 return None 

446 

447 def extend(self, stamp_list): 

448 """Extend Stamps instance by appending elements from another instance. 

449 

450 Parameters 

451 ---------- 

452 stamps_list : `list` [`Stamp`] 

453 List of Stamp object to append. 

454 """ 

455 for s in stamp_list: 

456 if not isinstance(s, Stamp): 

457 raise ValueError("Can only extend with Stamp objects") 

458 self._stamps += stamp_list 

459 

460 @classmethod 

461 def readFits(cls, filename): 

462 """Build an instance of this class from a file. 

463 

464 Parameters 

465 ---------- 

466 filename : `str` 

467 Name of the file to read. 

468 

469 Returns 

470 ------- 

471 object : `Stamps` 

472 An instance of this class. 

473 """ 

474 return cls.readFitsWithOptions(filename, None) 

475 

476 @classmethod 

477 def readFitsWithOptions(cls, filename, options): 

478 """Build an instance of this class with options. 

479 

480 Parameters 

481 ---------- 

482 filename : `str` 

483 Name of the file to read. 

484 options : `PropertyList` or `dict` 

485 Collection of metadata parameters. 

486 

487 Returns 

488 ------- 

489 object : `Stamps` 

490 An instance of this class. 

491 """ 

492 stamps, metadata = readFitsWithOptions(filename, Stamp.factory, options) 

493 return cls( 

494 stamps, 

495 metadata=metadata, 

496 use_mask=metadata["HAS_MASK"], 

497 use_variance=metadata["HAS_VARIANCE"], 

498 use_archive=metadata["HAS_ARCHIVE"], 

499 )