Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

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 

31 

32import lsst.afw.image as afwImage 

33import lsst.afw.fits as afwFits 

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

35from lsst.daf.base import PropertyList 

36from lsst.daf.butler.core.utils import getFullTypeName 

37from lsst.utils import doImport 

38 

39 

40def writeFits(filename, stamp_ims, metadata, type_name, write_mask, write_variance): 

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_ims : iterable of `lsst.afw.image.MaskedImageF` 

48 An iterable of masked images 

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

59 metadata['HAS_MASK'] = write_mask 

60 metadata['HAS_VARIANCE'] = write_variance 

61 metadata['N_STAMPS'] = len(stamp_ims) 

62 metadata['STAMPCLS'] = type_name 

63 # Record version number in case of future code changes 

64 metadata['VERSION'] = 1 

65 # create primary HDU with global metadata 

66 fitsPrimary = afwFits.Fits(filename, "w") 

67 fitsPrimary.createEmpty() 

68 fitsPrimary.writeMetadata(metadata) 

69 fitsPrimary.closeFile() 

70 

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

72 for i, stamp in enumerate(stamp_ims): 

73 metadata = PropertyList() 

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

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

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

77 if write_mask: 

78 metadata = PropertyList() 

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

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

81 if write_variance: 

82 metadata = PropertyList() 

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

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

85 return None 

86 

87 

88def readFitsWithOptions(filename, stamp_factory, options): 

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

90 subregion of the stamps to be read. 

91 

92 Parameters 

93 ---------- 

94 filename : `str` 

95 A string indicating the file to read 

96 stamp_factory : classmethod 

97 A factory function defined on a dataclass for constructing 

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

99 options : `PropertyList` or `dict` 

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

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

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

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

104 return a sub-image. 

105 

106 Returns 

107 ------- 

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

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

110 metadata : `PropertyList` 

111 The metadata 

112 """ 

113 # extract necessary info from metadata 

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

115 f = afwFits.Fits(filename, 'r') 

116 nExtensions = f.countHdus() 

117 nStamps = metadata["N_STAMPS"] 

118 # check if a bbox was provided 

119 kwargs = {} 

120 if options: 

121 # gen3 API 

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

123 kwargs["bbox"] = options["bbox"] 

124 # gen2 API 

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

126 llcX = options["llcX"] 

127 llcY = options["llcY"] 

128 width = options["width"] 

129 height = options["height"] 

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

131 kwargs["bbox"] = bbox 

132 stamp_parts = {} 

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

134 # header data unit 

135 for idx in range(nExtensions-1): 

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

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

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

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

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

141 else: 

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

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

144 if len(stamp_parts) != nStamps: 

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

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

147 # construct stamps themselves 

148 stamps = [] 

149 for k in range(nStamps): 

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

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

152 stamps.append(stamp_factory(maskedImage, metadata, k)) 

153 

154 return stamps, metadata 

155 

156 

157@dataclass 

158class AbstractStamp(abc.ABC): 

159 """Single abstract stamp 

160 

161 Parameters 

162 ---------- 

163 Inherit from this class to add metadata to the stamp 

164 """ 

165 

166 @classmethod 

167 @abc.abstractmethod 

168 def factory(cls, stamp_im, metadata, index): 

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

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

171 Parameters needed to construct this object are passed in via 

172 a metadata dictionary and then passed to the constructor of 

173 this class. 

174 

175 Parameters 

176 ---------- 

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

178 Pixel data to pass to the constructor 

179 metadata : `dict` 

180 Dictionary containing the information 

181 needed by the constructor. 

182 idx : `int` 

183 Index into the lists in ``metadata`` 

184 

185 Returns 

186 ------- 

187 stamp : `AbstractStamp` 

188 An instance of this class 

189 """ 

190 raise NotImplementedError 

191 

192 

193@dataclass 

194class Stamp(AbstractStamp): 

195 """Single stamp 

196 

197 Parameters 

198 ---------- 

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

200 The actual pixel values for the postage stamp 

201 position : `lsst.geom.SpherePoint` 

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

203 must keep track of the coordinate system 

204 """ 

205 stamp_im: afwImage.maskedImage.MaskedImageF 

206 position: SpherePoint 

207 

208 @classmethod 

209 def factory(cls, stamp_im, metadata, index): 

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

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

212 Parameters needed to construct this object are passed in via 

213 a metadata dictionary and then passed to the constructor of 

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

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

216 values will be passed: RA_DEG, DEC_DEG. They should 

217 each point to lists of values. 

218 

219 Parameters 

220 ---------- 

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

222 Pixel data to pass to the constructor 

223 metadata : `dict` 

224 Dictionary containing the information 

225 needed by the constructor. 

226 idx : `int` 

227 Index into the lists in ``metadata`` 

228 

229 Returns 

230 ------- 

231 stamp : `Stamp` 

232 An instance of this class 

233 """ 

234 if 'RA_DEG' in metadata and 'DEC_DEG' in metadata: 

235 return cls(stamp_im=stamp_im, 

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

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

238 else: 

239 return cls(stamp_im=stamp_im, position=SpherePoint(Angle(numpy.nan), Angle(numpy.nan))) 

240 

241 

242class StampsBase(abc.ABC, Sequence): 

243 """Collection of stamps and associated metadata. 

244 

245 Parameters 

246 ---------- 

247 stamps : iterable 

248 This should be an iterable of dataclass objects 

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

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

251 Metadata associated with the bright stars. 

252 use_mask : `bool`, optional 

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

254 use_variance : `bool`, optional 

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

256 

257 Notes 

258 ----- 

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

260 specified by a bbox: 

261 

262 >>> starSubregions = butler.get("brightStarStamps_sub", dataId, bbox=bbox) 

263 """ 

264 

265 def __init__(self, stamps, metadata=None, use_mask=True, use_variance=True): 

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

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

268 for stamp in stamps: 

269 if not isinstance(stamp, AbstractStamp): 

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

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

272 self._stamps = stamps 

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

274 self.use_mask = use_mask 

275 self.use_variance = use_variance 

276 

277 @classmethod 

278 def readFits(cls, filename): 

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

280 

281 Parameters 

282 ---------- 

283 filename : `str` 

284 Name of the file to read 

285 """ 

286 

287 return cls.readFitsWithOptions(filename, None) 

288 

289 @classmethod 

290 def readFitsWithOptions(cls, filename, options): 

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

292 

293 Parameters 

294 ---------- 

295 filename : `str` 

296 Name of the file to read 

297 options : `PropertyList` 

298 Collection of metadata parameters 

299 """ 

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

301 if cls is not StampsBase: 

302 raise NotImplementedError( 

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

304 ) 

305 

306 # Load metadata to get class 

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

308 type_name = metadata.get("STAMPCLS") 

309 if type_name is None: 

310 raise RuntimeError( 

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

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

313 ) 

314 

315 # Import class and override `cls` 

316 stamp_type = doImport(type_name) 

317 cls = stamp_type 

318 

319 return cls.readFitsWithOptions(filename, options) 

320 

321 @abc.abstractmethod 

322 def _refresh_metadata(self): 

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

324 can be extende 

325 """ 

326 raise NotImplementedError 

327 

328 def writeFits(self, filename): 

329 """Write this object to a file. 

330 

331 Parameters 

332 ---------- 

333 filename : `str` 

334 Name of file to write 

335 """ 

336 self._refresh_metadata() 

337 stamps_ims = self.getMaskedImages() 

338 type_name = getFullTypeName(self) 

339 writeFits(filename, stamps_ims, self._metadata, type_name, self.use_mask, self.use_variance) 

340 

341 def __len__(self): 

342 return len(self._stamps) 

343 

344 def __getitem__(self, index): 

345 return self._stamps[index] 

346 

347 def __iter__(self): 

348 return iter(self._stamps) 

349 

350 def getMaskedImages(self): 

351 """Retrieve star images. 

352 

353 Returns 

354 ------- 

355 maskedImages : 

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

357 """ 

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

359 

360 @property 

361 def metadata(self): 

362 return self._metadata 

363 

364 

365class Stamps(StampsBase): 

366 def _refresh_metadata(self): 

367 positions = self.getPositions() 

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

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

370 

371 def getPositions(self): 

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

373 

374 def append(self, item): 

375 """Add an additional stamp. 

376 

377 Parameters 

378 ---------- 

379 item : `Stamp` 

380 Stamp object to append. 

381 """ 

382 if not isinstance(item, Stamp): 

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

384 self._stamps.append(item) 

385 return None 

386 

387 def extend(self, stamp_list): 

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

389 

390 Parameters 

391 ---------- 

392 stamps_list : `list` [`Stamp`] 

393 List of Stamp object to append. 

394 """ 

395 for s in stamp_list: 

396 if not isinstance(s, Stamp): 

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

398 self._stamps += stamp_list 

399 

400 @classmethod 

401 def readFits(cls, filename): 

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

403 

404 Parameters 

405 ---------- 

406 filename : `str` 

407 Name of the file to read 

408 

409 Returns 

410 ------- 

411 object : `Stamps` 

412 An instance of this class 

413 """ 

414 return cls.readFitsWithOptions(filename, None) 

415 

416 @classmethod 

417 def readFitsWithOptions(cls, filename, options): 

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

419 

420 Parameters 

421 ---------- 

422 filename : `str` 

423 Name of the file to read 

424 options : `PropertyList` or `dict` 

425 Collection of metadata parameters 

426 

427 Returns 

428 ------- 

429 object : `Stamps` 

430 An instance of this class 

431 """ 

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

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

434 use_variance=metadata['HAS_VARIANCE'])