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

172 statements  

« prev     ^ index     » next       coverage.py v6.4.1, created at 2022-06-07 01:50 -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 

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

26 

27from collections.abc import Sequence 

28import abc 

29from dataclasses import dataclass 

30import numpy 

31from typing import Optional 

32 

33import lsst.afw.image as afwImage 

34import lsst.afw.fits as afwFits 

35import lsst.afw.table as afwTable 

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

37from lsst.daf.base import PropertyList 

38from lsst.utils.introspection import get_full_type_name 

39from lsst.utils import doImport 

40 

41 

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

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

44 

45 Parameters 

46 ---------- 

47 filename : `str` 

48 A string indicating the output filename 

49 stamps : iterable of `BaseStamp` 

50 An iterable of Stamp objects 

51 metadata : `PropertyList` 

52 A collection of key, value metadata pairs to be 

53 written to the primary header 

54 type_name : `str` 

55 Python type name of the StampsBase subclass to use 

56 write_mask : `bool` 

57 Write the mask data to the output file? 

58 write_variance : `bool` 

59 Write the variance data to the output file? 

60 write_archive : `bool`, optional 

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

62 Default: ``False``. 

63 """ 

64 metadata['HAS_MASK'] = write_mask 

65 metadata['HAS_VARIANCE'] = write_variance 

66 metadata['HAS_ARCHIVE'] = write_archive 

67 metadata['N_STAMPS'] = len(stamps) 

68 metadata['STAMPCLS'] = type_name 

69 # Record version number in case of future code changes 

70 metadata['VERSION'] = 1 

71 # create primary HDU with global metadata 

72 fitsFile = afwFits.Fits(filename, "w") 

73 fitsFile.createEmpty() 

74 # Store Persistables in an OutputArchive and write it 

75 if write_archive: 

76 oa = afwTable.io.OutputArchive() 

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

78 metadata["ARCHIVE_IDS"] = archive_ids 

79 fitsFile.writeMetadata(metadata) 

80 oa.writeFits(fitsFile) 

81 else: 

82 fitsFile.writeMetadata(metadata) 

83 fitsFile.closeFile() 

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

85 for i, stamp in enumerate(stamps): 

86 metadata = PropertyList() 

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

88 metadata.update({'EXTVER': i+1, 'EXTNAME': 'IMAGE'}) 

89 stamp.stamp_im.getImage().writeFits(filename, metadata=metadata, mode='a') 

90 if write_mask: 

91 metadata = PropertyList() 

92 metadata.update({'EXTVER': i+1, 'EXTNAME': 'MASK'}) 

93 stamp.stamp_im.getMask().writeFits(filename, metadata=metadata, mode='a') 

94 if write_variance: 

95 metadata = PropertyList() 

96 metadata.update({'EXTVER': i+1, 'EXTNAME': 'VARIANCE'}) 

97 stamp.stamp_im.getVariance().writeFits(filename, metadata=metadata, mode='a') 

98 return None 

99 

100 

101def readFitsWithOptions(filename, stamp_factory, options): 

102 """Read stamps from FITS file, allowing for only a 

103 subregion of the stamps to be read. 

104 

105 Parameters 

106 ---------- 

107 filename : `str` 

108 A string indicating the file to read 

109 stamp_factory : classmethod 

110 A factory function defined on a dataclass for constructing 

111 stamp objects a la `lsst.meas.alrogithm.Stamp` 

112 options : `PropertyList` or `dict` 

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

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

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

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

117 return a sub-image. 

118 

119 Returns 

120 ------- 

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

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

123 metadata : `PropertyList` 

124 The metadata 

125 """ 

126 # extract necessary info from metadata 

127 metadata = afwFits.readMetadata(filename, hdu=0) 

128 nStamps = metadata["N_STAMPS"] 

129 has_archive = metadata["HAS_ARCHIVE"] 

130 if has_archive: 

131 archive_ids = metadata.getArray("ARCHIVE_IDS") 

132 with afwFits.Fits(filename, 'r') as f: 

133 nExtensions = f.countHdus() 

134 # check if a bbox was provided 

135 kwargs = {} 

136 if options: 

137 # gen3 API 

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

139 kwargs["bbox"] = options["bbox"] 

140 # gen2 API 

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

142 llcX = options["llcX"] 

143 llcY = options["llcY"] 

144 width = options["width"] 

145 height = options["height"] 

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

147 kwargs["bbox"] = bbox 

148 stamp_parts = {} 

149 # We need to be careful because nExtensions includes the primary 

150 # header data unit 

151 for idx in range(nExtensions-1): 

152 md = afwFits.readMetadata(filename, hdu=idx+1) 

153 if md['EXTNAME'] in ('IMAGE', 'VARIANCE'): 

154 reader = afwImage.ImageFitsReader(filename, hdu=idx+1) 

155 elif md['EXTNAME'] == 'MASK': 

156 reader = afwImage.MaskFitsReader(filename, hdu=idx+1) 

157 elif md['EXTNAME'] == 'ARCHIVE_INDEX': 

158 f.setHdu(idx+1) 

159 archive = afwTable.io.InputArchive.readFits(f) 

160 continue 

161 elif md['EXTTYPE'] == 'ARCHIVE_DATA': 

162 continue 

163 else: 

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

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

166 if len(stamp_parts) != nStamps: 

167 raise ValueError(f'Number of stamps read ({len(stamp_parts)}) does not agree with the ' 

168 f'number of stamps recorded in the metadata ({nStamps}).') 

169 # construct stamps themselves 

170 stamps = [] 

171 for k in range(nStamps): 

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

173 maskedImage = afwImage.MaskedImageF(**stamp_parts[k+1]) 

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

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

176 

177 return stamps, metadata 

178 

179 

180@dataclass 

181class AbstractStamp(abc.ABC): 

182 """Single abstract stamp 

183 

184 Parameters 

185 ---------- 

186 Inherit from this class to add metadata to the stamp 

187 """ 

188 

189 @classmethod 

190 @abc.abstractmethod 

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

192 """This method is needed to service the FITS reader. 

193 We need a standard interface to construct objects like this. 

194 Parameters needed to construct this object are passed in via 

195 a metadata dictionary and then passed to the constructor of 

196 this class. 

197 

198 Parameters 

199 ---------- 

200 stamp : `lsst.afw.image.MaskedImage` 

201 Pixel data to pass to the constructor 

202 metadata : `dict` 

203 Dictionary containing the information 

204 needed by the constructor. 

205 idx : `int` 

206 Index into the lists in ``metadata`` 

207 archive_element : `lsst.afwTable.io.Persistable`, optional 

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

209 

210 Returns 

211 ------- 

212 stamp : `AbstractStamp` 

213 An instance of this class 

214 """ 

215 raise NotImplementedError 

216 

217 

218@dataclass 

219class Stamp(AbstractStamp): 

220 """Single stamp 

221 

222 Parameters 

223 ---------- 

224 stamp_im : `lsst.afw.image.MaskedImageF` 

225 The actual pixel values for the postage stamp 

226 position : `lsst.geom.SpherePoint`, optional 

227 Position of the center of the stamp. Note the user 

228 must keep track of the coordinate system 

229 """ 

230 stamp_im: afwImage.maskedImage.MaskedImageF 

231 archive_element: Optional[afwTable.io.Persistable] = None 

232 position: Optional[SpherePoint] = SpherePoint(Angle(numpy.nan), Angle(numpy.nan)) 

233 

234 @classmethod 

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

236 """This method is needed to service the FITS reader. 

237 We need a standard interface to construct objects like this. 

238 Parameters needed to construct this object are passed in via 

239 a metadata dictionary and then passed to the constructor of 

240 this class. If lists of values are passed with the following 

241 keys, they will be passed to the constructor, otherwise dummy 

242 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 : `afwTable.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(stamp_im=stamp_im, archive_element=archive_element, 

264 position=SpherePoint(Angle(metadata.getArray('RA_DEG')[index], degrees), 

265 Angle(metadata.getArray('DEC_DEG')[index], degrees))) 

266 else: 

267 return cls(stamp_im=stamp_im, archive_element=archive_element, 

268 position=SpherePoint(Angle(numpy.nan), Angle(numpy.nan))) 

269 

270 

271class StampsBase(abc.ABC, Sequence): 

272 """Collection of stamps and associated metadata. 

273 

274 Parameters 

275 ---------- 

276 stamps : iterable 

277 This should be an iterable of dataclass objects 

278 a la ``lsst.meas.algorithms.Stamp``. 

279 metadata : `lsst.daf.base.PropertyList`, optional 

280 Metadata associated with the objects within the stamps. 

281 use_mask : `bool`, optional 

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

283 use_variance : `bool`, optional 

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

285 use_archive : `bool`, optional 

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

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

288 Default ``False``. 

289 

290 Notes 

291 ----- 

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

293 specified by a bbox: 

294 

295 >>> starSubregions = butler.get("brightStarStamps", dataId, parameters={'bbox': bbox}) 

296 """ 

297 

298 def __init__(self, stamps, metadata=None, use_mask=True, use_variance=True, 

299 use_archive=False): 

300 if not hasattr(stamps, '__iter__'): 

301 raise ValueError('The stamps parameter must be iterable.') 

302 for stamp in stamps: 

303 if not isinstance(stamp, AbstractStamp): 

304 raise ValueError('The entries in stamps must inherit from AbstractStamp. ' 

305 f'Got {type(stamp)}.') 

306 self._stamps = stamps 

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

308 self.use_mask = use_mask 

309 self.use_variance = use_variance 

310 self.use_archive = use_archive 

311 

312 @classmethod 

313 def readFits(cls, filename): 

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

315 

316 Parameters 

317 ---------- 

318 filename : `str` 

319 Name of the file to read 

320 """ 

321 

322 return cls.readFitsWithOptions(filename, None) 

323 

324 @classmethod 

325 def readFitsWithOptions(cls, filename, options): 

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

327 

328 Parameters 

329 ---------- 

330 filename : `str` 

331 Name of the file to read 

332 options : `PropertyList` 

333 Collection of metadata parameters 

334 """ 

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

336 if cls is not StampsBase: 

337 raise NotImplementedError( 

338 f"Please implement specific FITS reader for class {cls}" 

339 ) 

340 

341 # Load metadata to get class 

342 metadata = afwFits.readMetadata(filename, hdu=0) 

343 type_name = metadata.get("STAMPCLS") 

344 if type_name is None: 

345 raise RuntimeError( 

346 f"No class name in file {filename}. Unable to instantiate correct" 

347 " stamps subclass. Is this an old version format Stamps file?" 

348 ) 

349 

350 # Import class and override `cls` 

351 stamp_type = doImport(type_name) 

352 cls = stamp_type 

353 

354 return cls.readFitsWithOptions(filename, options) 

355 

356 @abc.abstractmethod 

357 def _refresh_metadata(self): 

358 """Make sure metadata is up to date since this object 

359 can be extended 

360 """ 

361 raise NotImplementedError 

362 

363 def writeFits(self, filename): 

364 """Write this object to a file. 

365 

366 Parameters 

367 ---------- 

368 filename : `str` 

369 Name of file to write 

370 """ 

371 self._refresh_metadata() 

372 type_name = get_full_type_name(self) 

373 writeFits(filename, self._stamps, self._metadata, type_name, self.use_mask, self.use_variance, 

374 self.use_archive) 

375 

376 def __len__(self): 

377 return len(self._stamps) 

378 

379 def __getitem__(self, index): 

380 return self._stamps[index] 

381 

382 def __iter__(self): 

383 return iter(self._stamps) 

384 

385 def getMaskedImages(self): 

386 """Retrieve star images. 

387 

388 Returns 

389 ------- 

390 maskedImages : 

391 `list` [`lsst.afw.image.maskedImage.maskedImage.MaskedImageF`] 

392 """ 

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

394 

395 def getArchiveElements(self): 

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

397 

398 Returns 

399 ------- 

400 archiveElements : 

401 `list` [`lsst.afwTable.io.Persistable`] 

402 """ 

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

404 

405 @property 

406 def metadata(self): 

407 return self._metadata 

408 

409 

410class Stamps(StampsBase): 

411 def _refresh_metadata(self): 

412 positions = self.getPositions() 

413 self._metadata['RA_DEG'] = [p.getRa().asDegrees() for p in positions] 

414 self._metadata['DEC_DEG'] = [p.getDec().asDegrees() for p in positions] 

415 

416 def getPositions(self): 

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

418 

419 def append(self, item): 

420 """Add an additional stamp. 

421 

422 Parameters 

423 ---------- 

424 item : `Stamp` 

425 Stamp object to append. 

426 """ 

427 if not isinstance(item, Stamp): 

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

429 self._stamps.append(item) 

430 return None 

431 

432 def extend(self, stamp_list): 

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

434 

435 Parameters 

436 ---------- 

437 stamps_list : `list` [`Stamp`] 

438 List of Stamp object to append. 

439 """ 

440 for s in stamp_list: 

441 if not isinstance(s, Stamp): 

442 raise ValueError('Can only extend with Stamp objects') 

443 self._stamps += stamp_list 

444 

445 @classmethod 

446 def readFits(cls, filename): 

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

448 

449 Parameters 

450 ---------- 

451 filename : `str` 

452 Name of the file to read 

453 

454 Returns 

455 ------- 

456 object : `Stamps` 

457 An instance of this class 

458 """ 

459 return cls.readFitsWithOptions(filename, None) 

460 

461 @classmethod 

462 def readFitsWithOptions(cls, filename, options): 

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

464 

465 Parameters 

466 ---------- 

467 filename : `str` 

468 Name of the file to read 

469 options : `PropertyList` or `dict` 

470 Collection of metadata parameters 

471 

472 Returns 

473 ------- 

474 object : `Stamps` 

475 An instance of this class 

476 """ 

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

478 return cls(stamps, metadata=metadata, use_mask=metadata['HAS_MASK'], 

479 use_variance=metadata['HAS_VARIANCE'], use_archive=metadata['HAS_ARCHIVE'])