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 

36 

37 

38def writeFits(filename, stamp_ims, metadata, write_mask, write_variance): 

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

40 

41 Parameters 

42 ---------- 

43 filename : `str` 

44 A string indicating the output filename 

45 stamps_ims : iterable of `lsst.afw.image.MaskedImageF` 

46 An iterable of masked images 

47 metadata : `PropertyList` 

48 A collection of key, value metadata pairs to be 

49 written to the primary header 

50 write_mask : `bool` 

51 Write the mask data to the output file? 

52 write_variance : `bool` 

53 Write the variance data to the output file? 

54 """ 

55 metadata['HAS_MASK'] = write_mask 

56 metadata['HAS_VARIANCE'] = write_variance 

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

58 # create primary HDU with global metadata 

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

60 fitsPrimary.createEmpty() 

61 fitsPrimary.writeMetadata(metadata) 

62 fitsPrimary.closeFile() 

63 

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

65 for i, stamp in enumerate(stamp_ims): 

66 metadata = PropertyList() 

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

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

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

70 if write_mask: 

71 metadata = PropertyList() 

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

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

74 if write_variance: 

75 metadata = PropertyList() 

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

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

78 return None 

79 

80 

81def readFitsWithOptions(filename, stamp_factory, options): 

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

83 subregion of the stamps to be read. 

84 

85 Parameters 

86 ---------- 

87 filename : `str` 

88 A string indicating the file to read 

89 stamp_factory : classmethod 

90 A factory function defined on a dataclass for constructing 

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

92 options : `PropertyList` or `dict` 

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

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

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

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

97 return a sub-image. 

98 

99 Returns 

100 ------- 

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

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

103 metadata : `PropertyList` 

104 The metadata 

105 """ 

106 # extract necessary info from metadata 

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

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

109 nExtensions = f.countHdus() 

110 nStamps = metadata["N_STAMPS"] 

111 # check if a bbox was provided 

112 kwargs = {} 

113 if options: 

114 # gen3 API 

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

116 kwargs["bbox"] = options["bbox"] 

117 # gen2 API 

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

119 llcX = options["llcX"] 

120 llcY = options["llcY"] 

121 width = options["width"] 

122 height = options["height"] 

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

124 kwargs["bbox"] = bbox 

125 stamp_parts = {} 

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

127 # header data unit 

128 for idx in range(nExtensions-1): 

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

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

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

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

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

134 else: 

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

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

137 if len(stamp_parts) != nStamps: 

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

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

140 # construct stamps themselves 

141 stamps = [] 

142 for k in range(nStamps): 

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

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

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

146 

147 return stamps, metadata 

148 

149 

150@dataclass 

151class AbstractStamp(abc.ABC): 

152 """Single abstract stamp 

153 

154 Parameters 

155 ---------- 

156 Inherit from this class to add metadata to the stamp 

157 """ 

158 

159 @classmethod 

160 @abc.abstractmethod 

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

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

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

164 Parameters needed to construct this object are passed in via 

165 a metadata dictionary and then passed to the constructor of 

166 this class. 

167 

168 Parameters 

169 ---------- 

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

171 Pixel data to pass to the constructor 

172 metadata : `dict` 

173 Dictionary containing the information 

174 needed by the constructor. 

175 idx : `int` 

176 Index into the lists in ``metadata`` 

177 

178 Returns 

179 ------- 

180 stamp : `AbstractStamp` 

181 An instance of this class 

182 """ 

183 raise NotImplementedError 

184 

185 

186@dataclass 

187class Stamp(AbstractStamp): 

188 """Single stamp 

189 

190 Parameters 

191 ---------- 

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

193 The actual pixel values for the postage stamp 

194 position : `lsst.geom.SpherePoint` 

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

196 must keep track of the coordinate system 

197 """ 

198 stamp_im: afwImage.maskedImage.MaskedImageF 

199 position: SpherePoint 

200 

201 @classmethod 

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

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

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

205 Parameters needed to construct this object are passed in via 

206 a metadata dictionary and then passed to the constructor of 

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

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

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

210 each point to lists of values. 

211 

212 Parameters 

213 ---------- 

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

215 Pixel data to pass to the constructor 

216 metadata : `dict` 

217 Dictionary containing the information 

218 needed by the constructor. 

219 idx : `int` 

220 Index into the lists in ``metadata`` 

221 

222 Returns 

223 ------- 

224 stamp : `Stamp` 

225 An instance of this class 

226 """ 

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

228 return cls(stamp_im=stamp_im, 

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

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

231 else: 

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

233 

234 

235class StampsBase(abc.ABC, Sequence): 

236 """Collection of stamps and associated metadata. 

237 

238 Parameters 

239 ---------- 

240 stamps : iterable 

241 This should be an iterable of dataclass objects 

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

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

244 Metadata associated with the bright stars. 

245 use_mask : `bool`, optional 

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

247 use_variance : `bool`, optional 

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

249 

250 Notes 

251 ----- 

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

253 specified by a bbox: 

254 

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

256 """ 

257 

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

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

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

261 for stamp in stamps: 

262 if not isinstance(stamp, AbstractStamp): 

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

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

265 self._stamps = stamps 

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

267 self.use_mask = use_mask 

268 self.use_variance = use_variance 

269 

270 @classmethod 

271 @abc.abstractmethod 

272 def readFits(cls, filename): 

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

274 

275 Parameters 

276 ---------- 

277 filename : `str` 

278 Name of the file to read 

279 """ 

280 raise NotImplementedError 

281 

282 @classmethod 

283 @abc.abstractmethod 

284 def readFitsWithOptions(cls, filename, options): 

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

286 

287 Parameters 

288 ---------- 

289 filename : `str` 

290 Name of the file to read 

291 options : `PropertyList` 

292 Collection of metadata parameters 

293 """ 

294 raise NotImplementedError 

295 

296 @abc.abstractmethod 

297 def _refresh_metadata(self): 

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

299 can be extende 

300 """ 

301 raise NotImplementedError 

302 

303 def writeFits(self, filename): 

304 """Write this object to a file. 

305 

306 Parameters 

307 ---------- 

308 filename : `str` 

309 Name of file to write 

310 """ 

311 self._refresh_metadata() 

312 stamps_ims = self.getMaskedImages() 

313 writeFits(filename, stamps_ims, self._metadata, self.use_mask, self.use_variance) 

314 

315 def __len__(self): 

316 return len(self._stamps) 

317 

318 def __getitem__(self, index): 

319 return self._stamps[index] 

320 

321 def __iter__(self): 

322 return iter(self._stamps) 

323 

324 def getMaskedImages(self): 

325 """Retrieve star images. 

326 

327 Returns 

328 ------- 

329 maskedImages : 

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

331 """ 

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

333 

334 @property 

335 def metadata(self): 

336 return self._metadata 

337 

338 

339class Stamps(StampsBase): 

340 def _refresh_metadata(self): 

341 positions = self.getPositions() 

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

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

344 

345 def getPositions(self): 

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

347 

348 def append(self, item): 

349 """Add an additional stamp. 

350 

351 Parameters 

352 ---------- 

353 item : `Stamp` 

354 Stamp object to append. 

355 """ 

356 if not isinstance(item, Stamp): 

357 raise ValueError("Ojbects added must be a Stamp object.") 

358 self._stamps.append(item) 

359 return None 

360 

361 def extend(self, stamp_list): 

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

363 

364 Parameters 

365 ---------- 

366 stamps_list : `list` [`Stamp`] 

367 List of Stamp object to append. 

368 """ 

369 for s in stamp_list: 

370 if not isinstance(s, Stamp): 

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

372 self._stamps += stamp_list 

373 

374 @classmethod 

375 def readFits(cls, filename): 

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

377 

378 Parameters 

379 ---------- 

380 filename : `str` 

381 Name of the file to read 

382 

383 Returns 

384 ------- 

385 object : `Stamps` 

386 An instance of this class 

387 """ 

388 return cls.readFitsWithOptions(filename, None) 

389 

390 @classmethod 

391 def readFitsWithOptions(cls, filename, options): 

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

393 

394 Parameters 

395 ---------- 

396 filename : `str` 

397 Name of the file to read 

398 options : `PropertyList` or `dict` 

399 Collection of metadata parameters 

400 

401 Returns 

402 ------- 

403 object : `Stamps` 

404 An instance of this class 

405 """ 

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

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

408 use_variance=metadata['HAS_VARIANCE'])