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 

31from typing import Optional 

32 

33import lsst.afw.image as afwImage 

34import lsst.afw.fits as afwFits 

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

36from lsst.daf.base import PropertyList 

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

38from lsst.utils import doImport 

39 

40 

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

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

43 

44 Parameters 

45 ---------- 

46 filename : `str` 

47 A string indicating the output filename 

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

49 An iterable of masked images 

50 metadata : `PropertyList` 

51 A collection of key, value metadata pairs to be 

52 written to the primary header 

53 type_name : `str` 

54 Python type name of the StampsBase subclass to use 

55 write_mask : `bool` 

56 Write the mask data to the output file? 

57 write_variance : `bool` 

58 Write the variance data to the output file? 

59 """ 

60 metadata['HAS_MASK'] = write_mask 

61 metadata['HAS_VARIANCE'] = write_variance 

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

63 metadata['STAMPCLS'] = type_name 

64 # Record version number in case of future code changes 

65 metadata['VERSION'] = 1 

66 # create primary HDU with global metadata 

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

68 fitsPrimary.createEmpty() 

69 fitsPrimary.writeMetadata(metadata) 

70 fitsPrimary.closeFile() 

71 

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

73 for i, stamp in enumerate(stamp_ims): 

74 metadata = PropertyList() 

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

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

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

78 if write_mask: 

79 metadata = PropertyList() 

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

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

82 if write_variance: 

83 metadata = PropertyList() 

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

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

86 return None 

87 

88 

89def readFitsWithOptions(filename, stamp_factory, options): 

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

91 subregion of the stamps to be read. 

92 

93 Parameters 

94 ---------- 

95 filename : `str` 

96 A string indicating the file to read 

97 stamp_factory : classmethod 

98 A factory function defined on a dataclass for constructing 

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

100 options : `PropertyList` or `dict` 

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

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

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

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

105 return a sub-image. 

106 

107 Returns 

108 ------- 

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

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

111 metadata : `PropertyList` 

112 The metadata 

113 """ 

114 # extract necessary info from metadata 

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

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

117 nExtensions = f.countHdus() 

118 nStamps = metadata["N_STAMPS"] 

119 # check if a bbox was provided 

120 kwargs = {} 

121 if options: 

122 # gen3 API 

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

124 kwargs["bbox"] = options["bbox"] 

125 # gen2 API 

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

127 llcX = options["llcX"] 

128 llcY = options["llcY"] 

129 width = options["width"] 

130 height = options["height"] 

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

132 kwargs["bbox"] = bbox 

133 stamp_parts = {} 

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

135 # header data unit 

136 for idx in range(nExtensions-1): 

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

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

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

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

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

142 else: 

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

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

145 if len(stamp_parts) != nStamps: 

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

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

148 # construct stamps themselves 

149 stamps = [] 

150 for k in range(nStamps): 

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

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

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

154 

155 return stamps, metadata 

156 

157 

158@dataclass 

159class AbstractStamp(abc.ABC): 

160 """Single abstract stamp 

161 

162 Parameters 

163 ---------- 

164 Inherit from this class to add metadata to the stamp 

165 """ 

166 

167 @classmethod 

168 @abc.abstractmethod 

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

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

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

172 Parameters needed to construct this object are passed in via 

173 a metadata dictionary and then passed to the constructor of 

174 this class. 

175 

176 Parameters 

177 ---------- 

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

179 Pixel data to pass to the constructor 

180 metadata : `dict` 

181 Dictionary containing the information 

182 needed by the constructor. 

183 idx : `int` 

184 Index into the lists in ``metadata`` 

185 

186 Returns 

187 ------- 

188 stamp : `AbstractStamp` 

189 An instance of this class 

190 """ 

191 raise NotImplementedError 

192 

193 

194@dataclass 

195class Stamp(AbstractStamp): 

196 """Single stamp 

197 

198 Parameters 

199 ---------- 

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

201 The actual pixel values for the postage stamp 

202 position : `lsst.geom.SpherePoint` 

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

204 must keep track of the coordinate system 

205 """ 

206 stamp_im: afwImage.maskedImage.MaskedImageF 

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

208 

209 @classmethod 

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

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

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

213 Parameters needed to construct this object are passed in via 

214 a metadata dictionary and then passed to the constructor of 

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

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

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

218 each point to lists of values. 

219 

220 Parameters 

221 ---------- 

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

223 Pixel data to pass to the constructor 

224 metadata : `dict` 

225 Dictionary containing the information 

226 needed by the constructor. 

227 idx : `int` 

228 Index into the lists in ``metadata`` 

229 

230 Returns 

231 ------- 

232 stamp : `Stamp` 

233 An instance of this class 

234 """ 

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

236 return cls(stamp_im=stamp_im, 

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

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

239 else: 

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

241 

242 

243class StampsBase(abc.ABC, Sequence): 

244 """Collection of stamps and associated metadata. 

245 

246 Parameters 

247 ---------- 

248 stamps : iterable 

249 This should be an iterable of dataclass objects 

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

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

252 Metadata associated with the bright stars. 

253 use_mask : `bool`, optional 

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

255 use_variance : `bool`, optional 

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

257 

258 Notes 

259 ----- 

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

261 specified by a bbox: 

262 

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

264 """ 

265 

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

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

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

269 for stamp in stamps: 

270 if not isinstance(stamp, AbstractStamp): 

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

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

273 self._stamps = stamps 

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

275 self.use_mask = use_mask 

276 self.use_variance = use_variance 

277 

278 @classmethod 

279 def readFits(cls, filename): 

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

281 

282 Parameters 

283 ---------- 

284 filename : `str` 

285 Name of the file to read 

286 """ 

287 

288 return cls.readFitsWithOptions(filename, None) 

289 

290 @classmethod 

291 def readFitsWithOptions(cls, filename, options): 

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

293 

294 Parameters 

295 ---------- 

296 filename : `str` 

297 Name of the file to read 

298 options : `PropertyList` 

299 Collection of metadata parameters 

300 """ 

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

302 if cls is not StampsBase: 

303 raise NotImplementedError( 

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

305 ) 

306 

307 # Load metadata to get class 

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

309 type_name = metadata.get("STAMPCLS") 

310 if type_name is None: 

311 raise RuntimeError( 

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

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

314 ) 

315 

316 # Import class and override `cls` 

317 stamp_type = doImport(type_name) 

318 cls = stamp_type 

319 

320 return cls.readFitsWithOptions(filename, options) 

321 

322 @abc.abstractmethod 

323 def _refresh_metadata(self): 

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

325 can be extende 

326 """ 

327 raise NotImplementedError 

328 

329 def writeFits(self, filename): 

330 """Write this object to a file. 

331 

332 Parameters 

333 ---------- 

334 filename : `str` 

335 Name of file to write 

336 """ 

337 self._refresh_metadata() 

338 stamps_ims = self.getMaskedImages() 

339 type_name = getFullTypeName(self) 

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

341 

342 def __len__(self): 

343 return len(self._stamps) 

344 

345 def __getitem__(self, index): 

346 return self._stamps[index] 

347 

348 def __iter__(self): 

349 return iter(self._stamps) 

350 

351 def getMaskedImages(self): 

352 """Retrieve star images. 

353 

354 Returns 

355 ------- 

356 maskedImages : 

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

358 """ 

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

360 

361 @property 

362 def metadata(self): 

363 return self._metadata 

364 

365 

366class Stamps(StampsBase): 

367 def _refresh_metadata(self): 

368 positions = self.getPositions() 

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

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

371 

372 def getPositions(self): 

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

374 

375 def append(self, item): 

376 """Add an additional stamp. 

377 

378 Parameters 

379 ---------- 

380 item : `Stamp` 

381 Stamp object to append. 

382 """ 

383 if not isinstance(item, Stamp): 

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

385 self._stamps.append(item) 

386 return None 

387 

388 def extend(self, stamp_list): 

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

390 

391 Parameters 

392 ---------- 

393 stamps_list : `list` [`Stamp`] 

394 List of Stamp object to append. 

395 """ 

396 for s in stamp_list: 

397 if not isinstance(s, Stamp): 

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

399 self._stamps += stamp_list 

400 

401 @classmethod 

402 def readFits(cls, filename): 

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

404 

405 Parameters 

406 ---------- 

407 filename : `str` 

408 Name of the file to read 

409 

410 Returns 

411 ------- 

412 object : `Stamps` 

413 An instance of this class 

414 """ 

415 return cls.readFitsWithOptions(filename, None) 

416 

417 @classmethod 

418 def readFitsWithOptions(cls, filename, options): 

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

420 

421 Parameters 

422 ---------- 

423 filename : `str` 

424 Name of the file to read 

425 options : `PropertyList` or `dict` 

426 Collection of metadata parameters 

427 

428 Returns 

429 ------- 

430 object : `Stamps` 

431 An instance of this class 

432 """ 

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

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

435 use_variance=metadata['HAS_VARIANCE'])