Coverage for python/lsst/cell_coadds/_fits.py: 24%

95 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-03-28 03:59 -0700

1# This file is part of cell_coadds. 

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 

22from __future__ import annotations 

23 

24__all__ = ( 

25 "CellCoaddFitsFormatter", 

26 "CellCoaddFitsReader", 

27 "writeMultipleCellCoaddAsFits", 

28) 

29 

30import os 

31from collections.abc import Mapping 

32from typing import Any 

33 

34import lsst.afw.geom as afwGeom 

35import lsst.afw.image as afwImage 

36import numpy as np 

37from astropy.io import fits 

38from lsst.afw.image import ImageD, ImageF 

39from lsst.daf.base import PropertySet 

40from lsst.geom import Box2I, Extent2I, Point2I 

41from lsst.obs.base.formatters.fitsGeneric import FitsGenericFormatter 

42from lsst.skymap import Index2D 

43 

44from ._common_components import CoaddUnits, CommonComponents 

45from ._identifiers import CellIdentifiers, PatchIdentifiers 

46from ._image_planes import OwnedImagePlanes 

47from ._multiple_cell_coadd import MultipleCellCoadd, SingleCellCoadd 

48from ._uniform_grid import UniformGrid 

49 

50 

51class CellCoaddFitsFormatter(FitsGenericFormatter): 

52 """Interface for writing and reading cell coadds to/from FITS files. 

53 

54 This assumes the existence of readFits and writeFits methods (for now). 

55 """ 

56 

57 

58class CellCoaddFitsReader: 

59 """A reader class to read from a FITS file and produce cell-based coadds. 

60 

61 This reader class has read methods that can either return a single 

62 component without reading the entire file (e.g., readBBox, readWcs) 

63 and read methods that return a full coadd (e.g., 

64 readAsMultipleCellCoadd, readAsExplodedCellCoadd, readAsStitchedCoadd). 

65 

66 Parameters 

67 ---------- 

68 filename : `str` 

69 The name of the FITS file to read. 

70 """ 

71 

72 def __init__(self, filename: str) -> None: 

73 if not os.path.exists(filename): 

74 raise FileNotFoundError(f"File {filename} not found") 

75 

76 self.filename = filename 

77 

78 def readAsMultipleCellCoadd(self) -> MultipleCellCoadd: 

79 """Read the FITS file as a MultipleCellCoadd object.""" 

80 with fits.open(self.filename) as hdu_list: 

81 data = hdu_list[1].data 

82 header = hdu_list[1].header 

83 

84 # Read in WCS 

85 ps = PropertySet() 

86 ps.update(hdu_list[0].header) 

87 wcs = afwGeom.makeSkyWcs(ps) 

88 

89 # Build the quantities needed to construct a MultipleCellCoadd. 

90 common = CommonComponents( 

91 units=CoaddUnits(1), # TODO: read from FITS TUNIT1 (DM-40562) 

92 wcs=wcs, 

93 band=header["BAND"], 

94 identifiers=PatchIdentifiers( 

95 skymap=header["SKYMAP"], 

96 tract=header["TRACT"], 

97 patch=Index2D(x=header["PATCH_X"], y=header["PATCH_Y"]), 

98 band=header["BAND"], 

99 ), 

100 ) 

101 

102 grid_cell_size = Extent2I(header["GRCELL1"], header["GRCELL2"]) # Inner size of a single cell. 

103 grid_shape = Extent2I(header["GRSHAPE1"], header["GRSHAPE2"]) 

104 grid_min = Point2I(header["GRMIN1"], header["GRMIN2"]) 

105 grid = UniformGrid(cell_size=grid_cell_size, shape=grid_shape, min=grid_min) 

106 

107 # This is the inner bounding box for the multiple cell coadd 

108 inner_bbox = Box2I( 

109 Point2I(header["INBBOX11"], header["INBBOX12"]), 

110 Point2I(header["INBBOX21"], header["INBBOX22"]), 

111 ) 

112 

113 outer_cell_size = Extent2I(header["OCELL1"], header["OCELL2"]) 

114 psf_image_size = Extent2I(header["PSFSIZE1"], header["PSFSIZE2"]) 

115 

116 coadd = MultipleCellCoadd( 

117 ( 

118 self._readSingleCellCoadd( 

119 data=row, 

120 header=header, 

121 common=common, 

122 outer_cell_size=outer_cell_size, 

123 psf_image_size=psf_image_size, 

124 inner_cell_size=grid_cell_size, 

125 ) 

126 for row in data 

127 ), 

128 grid=grid, 

129 outer_cell_size=outer_cell_size, 

130 psf_image_size=psf_image_size, 

131 inner_bbox=inner_bbox, 

132 common=common, 

133 ) 

134 

135 return coadd 

136 

137 @staticmethod 

138 def _readSingleCellCoadd( 

139 data: Mapping[str, Any], 

140 common: CommonComponents, 

141 header: Mapping[str, Any], 

142 *, 

143 outer_cell_size: Extent2I, 

144 inner_cell_size: Extent2I, 

145 psf_image_size: Extent2I, 

146 ) -> SingleCellCoadd: 

147 """Read a coadd from a FITS file. 

148 

149 Parameters 

150 ---------- 

151 data : `Mapping` 

152 The data from the FITS file. Usually, a single row from the binary 

153 table representation. 

154 common : `CommonComponents` 

155 The common components of the coadd. 

156 outer_cell_size : `Extent2I` 

157 The size of the outer cell. 

158 psf_image_size : `Extent2I` 

159 The size of the PSF image. 

160 inner_cell_size : `Extent2I` 

161 The size of the inner cell. 

162 

163 Returns 

164 ------- 

165 coadd : `SingleCellCoadd` 

166 The coadd read from the file. 

167 """ 

168 buffer = (outer_cell_size - inner_cell_size) // 2 

169 

170 psf = ImageD( 

171 array=data["psf"].astype(np.float64), 

172 xy0=(-(psf_image_size // 2)).asPoint(), # integer division and negation do not commute. 

173 ) # use the variable 

174 xy0 = Point2I( 

175 inner_cell_size.x * data["cell_id"][0] - buffer.x + header["GRMIN1"], 

176 inner_cell_size.y * data["cell_id"][1] - buffer.y + header["GRMIN2"], 

177 ) 

178 mask = afwImage.Mask(data["mask"].astype(np.int32), xy0=xy0) 

179 image_planes = OwnedImagePlanes( 

180 image=ImageF( 

181 data["image"].astype(np.float32), 

182 xy0=xy0, 

183 ), 

184 mask=mask, 

185 variance=ImageF(data["variance"].astype(np.float32), xy0=xy0), 

186 noise_realizations=[], 

187 mask_fractions=None, 

188 ) 

189 

190 identifiers = CellIdentifiers( 

191 cell=Index2D(data["cell_id"][0], data["cell_id"][1]), 

192 skymap=common.identifiers.skymap, 

193 tract=common.identifiers.tract, 

194 patch=common.identifiers.patch, 

195 band=common.identifiers.band, 

196 ) 

197 

198 return SingleCellCoadd( 

199 outer=image_planes, 

200 psf=psf, 

201 inner_bbox=Box2I( 

202 corner=Point2I( 

203 inner_cell_size.x * data["cell_id"][0] + header["GRMIN1"], 

204 inner_cell_size.y * data["cell_id"][1] + header["GRMIN2"], 

205 ), 

206 dimensions=inner_cell_size, 

207 ), 

208 common=common, 

209 identifiers=identifiers, 

210 # TODO: Pass a sensible value here in DM-40563. 

211 inputs=None, # type: ignore[arg-type] 

212 ) 

213 

214 def readWcs(self) -> afwGeom.SkyWcs: 

215 """Read the WCS information from the FITS file. 

216 

217 Returns 

218 ------- 

219 wcs : `~lsst.afw.geom.SkyWcs` 

220 The WCS information read from the FITS file. 

221 """ 

222 # Read in WCS 

223 ps = PropertySet() 

224 with fits.open(self.filename) as hdu_list: 

225 ps.update(hdu_list[0].header) 

226 wcs = afwGeom.makeSkyWcs(ps) 

227 return wcs 

228 

229 

230def writeMultipleCellCoaddAsFits( 

231 multiple_cell_coadd: MultipleCellCoadd, 

232 filename: str, 

233 overwrite: bool = False, 

234 metadata: PropertySet | None = None, 

235) -> None: 

236 """Write a MultipleCellCoadd object to a FITS file. 

237 

238 Parameters 

239 ---------- 

240 multiple_cell_coadd : `MultipleCellCoadd` 

241 The multiple cell coadd to write to a FITS file. 

242 filename : `str` 

243 The name of the file to write to. 

244 overwrite : `bool`, optional 

245 Whether to overwrite the file if it already exists? 

246 metadata : `~lsst.daf.base.PropertySet`, optional 

247 Additional metadata to write to the FITS file. 

248 """ 

249 cell_id = fits.Column( 

250 name="cell_id", 

251 format="2I", 

252 array=[cell.identifiers.cell for cell in multiple_cell_coadd.cells.values()], 

253 ) 

254 

255 image_array = [cell.outer.image.array for cell in multiple_cell_coadd.cells.values()] 

256 unit_array = [cell.common.units.name for cell in multiple_cell_coadd.cells.values()] 

257 image = fits.Column( 

258 name="image", 

259 unit=unit_array[0], 

260 format=f"{image_array[0].size}E", 

261 dim=f"({image_array[0].shape[1]}, {image_array[0].shape[0]})", 

262 array=image_array, 

263 ) 

264 

265 mask_array = [cell.outer.mask.array for cell in multiple_cell_coadd.cells.values()] 

266 mask = fits.Column( 

267 name="mask", 

268 format=f"{mask_array[0].size}I", 

269 dim=f"({mask_array[0].shape[1]}, {mask_array[0].shape[0]})", 

270 array=mask_array, 

271 ) 

272 

273 variance_array = [cell.outer.variance.array for cell in multiple_cell_coadd.cells.values()] 

274 variance = fits.Column( 

275 name="variance", 

276 format=f"{variance_array[0].size}E", 

277 dim=f"({variance_array[0].shape[1]}, {variance_array[0].shape[0]})", 

278 array=variance_array, 

279 ) 

280 

281 psf_array = [cell.psf_image.array for cell in multiple_cell_coadd.cells.values()] 

282 psf = fits.Column( 

283 name="psf", 

284 format=f"{psf_array[0].size}D", 

285 dim=f"({psf_array[0].shape[1]}, {psf_array[0].shape[0]})", 

286 array=[cell.psf_image.array for cell in multiple_cell_coadd.cells.values()], 

287 ) 

288 

289 col_defs = fits.ColDefs([cell_id, image, mask, variance, psf]) 

290 hdu = fits.BinTableHDU.from_columns(col_defs) 

291 

292 grid_cell_size = multiple_cell_coadd.grid.cell_size 

293 grid_shape = multiple_cell_coadd.grid.shape 

294 grid_min = multiple_cell_coadd.grid.bbox.getMin() 

295 grid_cards = { 

296 "GRCELL1": grid_cell_size.x, 

297 "GRCELL2": grid_cell_size.y, 

298 "GRSHAPE1": grid_shape.x, 

299 "GRSHAPE2": grid_shape.y, 

300 "GRMIN1": grid_min.x, 

301 "GRMIN2": grid_min.y, 

302 } 

303 hdu.header.extend(grid_cards) 

304 

305 outer_cell_size_cards = { 

306 "OCELL1": multiple_cell_coadd.outer_cell_size.x, 

307 "OCELL2": multiple_cell_coadd.outer_cell_size.y, 

308 } 

309 hdu.header.extend(outer_cell_size_cards) 

310 

311 psf_image_size_cards = { 

312 "PSFSIZE1": multiple_cell_coadd.psf_image_size.x, 

313 "PSFSIZE2": multiple_cell_coadd.psf_image_size.y, 

314 } 

315 hdu.header.extend(psf_image_size_cards) 

316 

317 inner_bbox_cards = { 

318 "INBBOX11": multiple_cell_coadd.inner_bbox.minX, 

319 "INBBOX12": multiple_cell_coadd.inner_bbox.minY, 

320 "INBBOX21": multiple_cell_coadd.inner_bbox.maxX, 

321 "INBBOX22": multiple_cell_coadd.inner_bbox.maxY, 

322 } 

323 hdu.header.extend(inner_bbox_cards) 

324 

325 wcs = multiple_cell_coadd.common.wcs 

326 wcs_cards = wcs.getFitsMetadata().toDict() 

327 primary_hdu = fits.PrimaryHDU() 

328 primary_hdu.header.extend(wcs_cards) 

329 

330 hdu.header["TUNIT1"] = multiple_cell_coadd.common.units.name 

331 # This assumed to be the same as multiple_cell_coadd.common.identifers.band 

332 # See DM-38843. 

333 hdu.header["BAND"] = multiple_cell_coadd.common.band 

334 hdu.header["SKYMAP"] = multiple_cell_coadd.common.identifiers.skymap 

335 hdu.header["TRACT"] = multiple_cell_coadd.common.identifiers.tract 

336 hdu.header["PATCH_X"] = multiple_cell_coadd.common.identifiers.patch.x 

337 hdu.header["PATCH_Y"] = multiple_cell_coadd.common.identifiers.patch.y 

338 

339 if metadata is not None: 

340 hdu.header.extend(metadata.toDict()) 

341 

342 hdu_list = fits.HDUList([primary_hdu, hdu]) 

343 hdu_list.writeto(filename, overwrite=overwrite)