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

183 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2023-10-17 10:00 +0000

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, field, fields 

29 

30import numpy as np 

31from lsst.afw.fits import Fits, readMetadata 

32from lsst.afw.image import ImageFitsReader, MaskedImage, 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.algorithm.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 Notes 

125 ----- 

126 The data are read using the data type expected by the 

127 `~lsst.afw.image.MaskedImage` class attached to the `AbstractStamp` 

128 dataclass associated with the factory method. 

129 """ 

130 # extract necessary info from metadata 

131 metadata = readMetadata(filename, hdu=0) 

132 nStamps = metadata["N_STAMPS"] 

133 has_archive = metadata["HAS_ARCHIVE"] 

134 if has_archive: 

135 archive_ids = metadata.getArray("ARCHIVE_IDS") 

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

137 nExtensions = f.countHdus() 

138 # check if a bbox was provided 

139 kwargs = {} 

140 if options: 

141 # gen3 API 

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

143 kwargs["bbox"] = options["bbox"] 

144 # gen2 API 

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

146 llcX = options["llcX"] 

147 llcY = options["llcY"] 

148 width = options["width"] 

149 height = options["height"] 

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

151 kwargs["bbox"] = bbox 

152 stamp_parts = {} 

153 

154 # Determine the dtype from the factory. This allows a Stamp class 

155 # to be defined in terms of MaskedImageD or MaskedImageI without 

156 # forcing everything to floats. 

157 masked_image_cls = None 

158 for stamp_field in fields(stamp_factory.__self__): 

159 if issubclass(stamp_field.type, MaskedImage): 

160 masked_image_cls = stamp_field.type 

161 break 

162 else: 

163 raise RuntimeError("Stamp factory does not use MaskedImage.") 

164 default_dtype = np.dtype(masked_image_cls.dtype) 

165 variance_dtype = np.dtype(np.float32) # Variance is always the same type. 

166 

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

168 for idx in range(nExtensions - 1): 

169 dtype = None 

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

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

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

173 if md["EXTNAME"] == "VARIANCE": 

174 dtype = variance_dtype 

175 else: 

176 dtype = default_dtype 

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

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

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

180 f.setHdu(idx + 1) 

181 archive = InputArchive.readFits(f) 

182 continue 

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

184 continue 

185 else: 

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

187 stamp_parts.setdefault(md["EXTVER"], {})[md["EXTNAME"].lower()] = reader.read(dtype=dtype, 

188 **kwargs) 

189 if len(stamp_parts) != nStamps: 

190 raise ValueError( 

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

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

193 ) 

194 # construct stamps themselves 

195 stamps = [] 

196 for k in range(nStamps): 

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

198 maskedImage = masked_image_cls(**stamp_parts[k + 1]) 

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

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

201 

202 return stamps, metadata 

203 

204 

205@dataclass 

206class AbstractStamp(abc.ABC): 

207 """Single abstract stamp. 

208 

209 Parameters 

210 ---------- 

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

212 """ 

213 

214 @classmethod 

215 @abc.abstractmethod 

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

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

218 interface to construct objects like this. Parameters needed to 

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

220 passed to the constructor of this class. 

221 

222 Parameters 

223 ---------- 

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

225 Pixel data to pass to the constructor 

226 metadata : `dict` 

227 Dictionary containing the information 

228 needed by the constructor. 

229 idx : `int` 

230 Index into the lists in ``metadata`` 

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

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

233 

234 Returns 

235 ------- 

236 stamp : `AbstractStamp` 

237 An instance of this class 

238 """ 

239 raise NotImplementedError 

240 

241 

242def _default_position(): 

243 # SpherePoint is nominally mutable in C++ so we must use a factory 

244 # and return an entirely new SpherePoint each time a Stamps is created. 

245 return SpherePoint(Angle(np.nan), Angle(np.nan)) 

246 

247 

248@dataclass 

249class Stamp(AbstractStamp): 

250 """Single stamp. 

251 

252 Parameters 

253 ---------- 

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

255 The actual pixel values for the postage stamp. 

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

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

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

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

260 the coordinate system. 

261 """ 

262 

263 stamp_im: MaskedImageF 

264 archive_element: Persistable | None = None 

265 position: SpherePoint | None = field(default_factory=_default_position) 

266 

267 @classmethod 

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

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

270 interface to construct objects like this. Parameters needed to 

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

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

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

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

275 each point to lists of values. 

276 

277 Parameters 

278 ---------- 

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

280 Pixel data to pass to the constructor 

281 metadata : `dict` 

282 Dictionary containing the information 

283 needed by the constructor. 

284 idx : `int` 

285 Index into the lists in ``metadata`` 

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

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

288 

289 Returns 

290 ------- 

291 stamp : `Stamp` 

292 An instance of this class 

293 """ 

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

295 return cls( 

296 stamp_im=stamp_im, 

297 archive_element=archive_element, 

298 position=SpherePoint( 

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

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

301 ), 

302 ) 

303 else: 

304 return cls( 

305 stamp_im=stamp_im, 

306 archive_element=archive_element, 

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

308 ) 

309 

310 

311class StampsBase(abc.ABC, Sequence): 

312 """Collection of stamps and associated metadata. 

313 

314 Parameters 

315 ---------- 

316 stamps : iterable 

317 This should be an iterable of dataclass objects 

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

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

320 Metadata associated with the objects within the stamps. 

321 use_mask : `bool`, optional 

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

323 use_variance : `bool`, optional 

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

325 use_archive : `bool`, optional 

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

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

328 Default ``False``. 

329 

330 Notes 

331 ----- 

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

333 specified by a bbox: 

334 

335 >>> starSubregions = butler.get( 

336 "brightStarStamps", 

337 dataId, 

338 parameters={"bbox": bbox} 

339 ) 

340 """ 

341 

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

343 for stamp in stamps: 

344 if not isinstance(stamp, AbstractStamp): 

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

346 self._stamps = stamps 

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

348 self.use_mask = use_mask 

349 self.use_variance = use_variance 

350 self.use_archive = use_archive 

351 

352 @classmethod 

353 def readFits(cls, filename): 

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

355 

356 Parameters 

357 ---------- 

358 filename : `str` 

359 Name of the file to read 

360 """ 

361 

362 return cls.readFitsWithOptions(filename, None) 

363 

364 @classmethod 

365 def readFitsWithOptions(cls, filename, options): 

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

367 

368 Parameters 

369 ---------- 

370 filename : `str` 

371 Name of the file to read 

372 options : `PropertyList` 

373 Collection of metadata parameters 

374 """ 

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

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

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

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

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

380 if cls is not StampsBase: 

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

382 

383 # Load metadata to get class 

384 metadata = readMetadata(filename, hdu=0) 

385 type_name = metadata.get("STAMPCLS") 

386 if type_name is None: 

387 raise RuntimeError( 

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

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

390 ) 

391 

392 # Import class and override `cls` 

393 stamp_type = doImport(type_name) 

394 cls = stamp_type 

395 

396 return cls.readFitsWithOptions(filename, options) 

397 

398 @abc.abstractmethod 

399 def _refresh_metadata(self): 

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

401 raise NotImplementedError 

402 

403 def writeFits(self, filename): 

404 """Write this object to a file. 

405 

406 Parameters 

407 ---------- 

408 filename : `str` 

409 Name of file to write. 

410 """ 

411 self._refresh_metadata() 

412 type_name = get_full_type_name(self) 

413 writeFits( 

414 filename, 

415 self._stamps, 

416 self._metadata, 

417 type_name, 

418 self.use_mask, 

419 self.use_variance, 

420 self.use_archive, 

421 ) 

422 

423 def __len__(self): 

424 return len(self._stamps) 

425 

426 def __getitem__(self, index): 

427 return self._stamps[index] 

428 

429 def __iter__(self): 

430 return iter(self._stamps) 

431 

432 def getMaskedImages(self): 

433 """Retrieve star images. 

434 

435 Returns 

436 ------- 

437 maskedImages : 

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

439 """ 

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

441 

442 def getArchiveElements(self): 

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

444 

445 Returns 

446 ------- 

447 archiveElements : 

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

449 """ 

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

451 

452 @property 

453 def metadata(self): 

454 return self._metadata 

455 

456 

457class Stamps(StampsBase): 

458 def _refresh_metadata(self): 

459 positions = self.getPositions() 

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

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

462 

463 def getPositions(self): 

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

465 

466 def append(self, item): 

467 """Add an additional stamp. 

468 

469 Parameters 

470 ---------- 

471 item : `Stamp` 

472 Stamp object to append. 

473 """ 

474 if not isinstance(item, Stamp): 

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

476 self._stamps.append(item) 

477 return None 

478 

479 def extend(self, stamp_list): 

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

481 

482 Parameters 

483 ---------- 

484 stamps_list : `list` [`Stamp`] 

485 List of Stamp object to append. 

486 """ 

487 for s in stamp_list: 

488 if not isinstance(s, Stamp): 

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

490 self._stamps += stamp_list 

491 

492 @classmethod 

493 def readFits(cls, filename): 

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

495 

496 Parameters 

497 ---------- 

498 filename : `str` 

499 Name of the file to read. 

500 

501 Returns 

502 ------- 

503 object : `Stamps` 

504 An instance of this class. 

505 """ 

506 return cls.readFitsWithOptions(filename, None) 

507 

508 @classmethod 

509 def readFitsWithOptions(cls, filename, options): 

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

511 

512 Parameters 

513 ---------- 

514 filename : `str` 

515 Name of the file to read. 

516 options : `PropertyList` or `dict` 

517 Collection of metadata parameters. 

518 

519 Returns 

520 ------- 

521 object : `Stamps` 

522 An instance of this class. 

523 """ 

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

525 return cls( 

526 stamps, 

527 metadata=metadata, 

528 use_mask=metadata["HAS_MASK"], 

529 use_variance=metadata["HAS_VARIANCE"], 

530 use_archive=metadata["HAS_ARCHIVE"], 

531 )