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

115 statements  

« prev     ^ index     » next       coverage.py v7.5.0, created at 2024-05-03 11:12 +0000

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 

22"""Module to handle FITS serialization and de-serialization. 

23 

24The routines to write and read the files are in the same module, as a change to 

25one is typically accompanied by a corresponding change to another. Code changes 

26relating to writing the file must bump to the version number denoted by the 

27module constant FILE_FORMAT_VERSION. 

28 

29Although the typical use case is for newer versions of the code to read files 

30written by an older version, for the purposes of deciding the newer version 

31string, it is helpful to think about an older version of the reader attempting 

32to read a newer version of the file on disk. The policy for bumping the version 

33is as follows: 

34 

351. When the on-disk file format written by this module changes such that the 

36previous version of the reader can still read files written by the newer 

37version, then there should be a minor bump. 

38 

392. When the on-disk format written by this module changes in a way that will 

40prevent the previous version of the reader from reading a file produced by the 

41current version of the module, then there should be a major bump. This usually 

42means that the new version of the reader cannot read older file either, 

43save the temporary support with deprecation warnings, possibly until a new 

44release of the Science Pipelines is made. 

45 

46Examples 

47-------- 

481. A file with VERSION=1.3 should still be readable by the reader in 

49this module when the module-level constant FILE_FORMAT_VERSION=1.4. A file 

50written with VERSION=1.4 will typically be readable by a reader when the 

51module-level FILE_FORMAT_VERSION=1.3, although such a use case is not expected. 

52A concrete example of change 

53that requires only a minor bump is adding another BinTable that keeps track of 

54the input visits. 

55 

562. An example of major change would be migrating from using 

57BinTableHDU to ImageHDU to save data. Even if the reader supports reading 

58either of this formats based on the value of VERSION from the header, it should 

59be a major change because the previous version of the reader cannot read data 

60from ImageHDUs. 

61 

62Unit tests only check that a file written can be read by the concurrent version 

63of the module, but not by any of the previous ones. Hence, bumping 

64FILE_FORMAT_VERSION to the appropriate value is ultimately at the discretion of 

65the developers. 

66 

67A major bump must also be recorded in the `isCompatibleWith` method. 

68It is plausible that different (non-consequent) major format versions can be 

69read by the same reader (due to reverting back to an earlier format, or to 

70something very similar). `isCompatibleWith` method offers the convenience of 

71checking if a particular format version can be read by the current reader. 

72 

73Note that major version 0 is considered unstable and experimental and none of 

74the guarantee above applies. 

75""" 

76 

77from __future__ import annotations 

78 

79__all__ = ( 

80 "CellCoaddFitsFormatter", 

81 "CellCoaddFitsReader", 

82 "IncompatibleVersionError", 

83 "writeMultipleCellCoaddAsFits", 

84) 

85 

86import logging 

87import os 

88from collections.abc import Mapping 

89from typing import Any 

90 

91import lsst.afw.geom as afwGeom 

92import lsst.afw.image as afwImage 

93import numpy as np 

94from astropy.io import fits 

95from lsst.afw.image import ImageD, ImageF 

96from lsst.daf.base import PropertySet 

97from lsst.geom import Box2I, Extent2I, Point2I 

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

99from lsst.skymap import Index2D 

100 

101from ._common_components import CoaddUnits, CommonComponents 

102from ._identifiers import CellIdentifiers, PatchIdentifiers 

103from ._image_planes import OwnedImagePlanes 

104from ._multiple_cell_coadd import MultipleCellCoadd, SingleCellCoadd 

105from ._uniform_grid import UniformGrid 

106 

107FILE_FORMAT_VERSION = "0.2" 

108"""Version number for the file format as persisted, presented as a string of 

109the form M.m, where M is the major version, m is the minor version. 

110""" 

111 

112logger = logging.getLogger(__name__) 

113 

114 

115class IncompatibleVersionError(RuntimeError): 

116 """Exception raised when the CellCoaddFitsReader version is not compatible 

117 with the FITS file attempted to read. 

118 """ 

119 

120 

121class CellCoaddFitsFormatter(FitsGenericFormatter): 

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

123 

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

125 """ 

126 

127 

128class CellCoaddFitsReader: 

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

130 

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

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

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

134 readAsMultipleCellCoadd, readAsExplodedCellCoadd, readAsStitchedCoadd). 

135 

136 Parameters 

137 ---------- 

138 filename : `str` 

139 The name of the FITS file to read. 

140 """ 

141 

142 # Minimum and maximum compatible file format versions are listed as 

143 # iterables so as to allow for discontiguous intervals. 

144 MINIMUM_FILE_FORMAT_VERSIONS = ("0.1",) 

145 MAXIMUM_FILE_FORMAT_VERSIONS = ("1.0",) 

146 

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

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

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

150 

151 self.filename = filename 

152 

153 @classmethod 

154 def isCompatibleWith(cls, written_version: str, /) -> bool: 

155 """Check if the serialization version is compatible with the reader. 

156 

157 This is a convenience method to ask if the current version of this 

158 class can read a file, based on the VERSION in its header. 

159 

160 Parameters 

161 ---------- 

162 written_version: `str` 

163 The VERSION of the file to be read. 

164 

165 Returns 

166 ------- 

167 compatible : `bool` 

168 Whether the reader can read a file whose VERSION is 

169 ``written_version``. 

170 

171 Notes 

172 ----- 

173 This accepts the other version as a positional argument only. 

174 """ 

175 for min_version, max_version in zip( 

176 cls.MINIMUM_FILE_FORMAT_VERSIONS, 

177 cls.MAXIMUM_FILE_FORMAT_VERSIONS, 

178 strict=True, 

179 ): 

180 if min_version <= written_version < max_version: 

181 return True 

182 

183 return False 

184 

185 def readAsMultipleCellCoadd(self) -> MultipleCellCoadd: 

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

187 

188 Raises 

189 ------ 

190 IncompatibleError 

191 Raised if the version of this module that wrote the file is 

192 incompatible with this module that is reading it in. 

193 """ 

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

195 header = hdu_list[1].header 

196 written_version = header.get("VERSION", "0.1") 

197 if not self.isCompatibleWith(written_version): 

198 raise IncompatibleVersionError( 

199 f"{self.filename} was written with version {written_version}" 

200 f"but attempting to read it with a reader designed for {FILE_FORMAT_VERSION}" 

201 ) 

202 if written_version != FILE_FORMAT_VERSION: 

203 logger.info( 

204 "Reading %s having version %s with reader designed for %s", 

205 self.filename, 

206 written_version, 

207 FILE_FORMAT_VERSION, 

208 ) 

209 

210 data = hdu_list[1].data 

211 

212 # Read in WCS 

213 ps = PropertySet() 

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

215 wcs = afwGeom.makeSkyWcs(ps) 

216 

217 # Build the quantities needed to construct a MultipleCellCoadd. 

218 common = CommonComponents( 

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

220 wcs=wcs, 

221 band=header["BAND"], 

222 identifiers=PatchIdentifiers( 

223 skymap=header["SKYMAP"], 

224 tract=header["TRACT"], 

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

226 band=header["BAND"], 

227 ), 

228 ) 

229 

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

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

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

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

234 

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

236 inner_bbox = Box2I( 

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

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

239 ) 

240 

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

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

243 

244 coadd = MultipleCellCoadd( 

245 ( 

246 self._readSingleCellCoadd( 

247 data=row, 

248 header=header, 

249 common=common, 

250 outer_cell_size=outer_cell_size, 

251 psf_image_size=psf_image_size, 

252 inner_cell_size=grid_cell_size, 

253 ) 

254 for row in data 

255 ), 

256 grid=grid, 

257 outer_cell_size=outer_cell_size, 

258 psf_image_size=psf_image_size, 

259 inner_bbox=inner_bbox, 

260 common=common, 

261 ) 

262 

263 return coadd 

264 

265 @staticmethod 

266 def _readSingleCellCoadd( 

267 data: Mapping[str, Any], 

268 common: CommonComponents, 

269 header: Mapping[str, Any], 

270 *, 

271 outer_cell_size: Extent2I, 

272 inner_cell_size: Extent2I, 

273 psf_image_size: Extent2I, 

274 ) -> SingleCellCoadd: 

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

276 

277 Parameters 

278 ---------- 

279 data : `Mapping` 

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

281 table representation. 

282 common : `CommonComponents` 

283 The common components of the coadd. 

284 outer_cell_size : `Extent2I` 

285 The size of the outer cell. 

286 psf_image_size : `Extent2I` 

287 The size of the PSF image. 

288 inner_cell_size : `Extent2I` 

289 The size of the inner cell. 

290 

291 Returns 

292 ------- 

293 coadd : `SingleCellCoadd` 

294 The coadd read from the file. 

295 """ 

296 buffer = (outer_cell_size - inner_cell_size) // 2 

297 

298 psf = ImageD( 

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

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

301 ) # use the variable 

302 xy0 = Point2I( 

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

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

305 ) 

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

307 image_planes = OwnedImagePlanes( 

308 image=ImageF( 

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

310 xy0=xy0, 

311 ), 

312 mask=mask, 

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

314 noise_realizations=[], 

315 mask_fractions=None, 

316 ) 

317 

318 identifiers = CellIdentifiers( 

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

320 skymap=common.identifiers.skymap, 

321 tract=common.identifiers.tract, 

322 patch=common.identifiers.patch, 

323 band=common.identifiers.band, 

324 ) 

325 

326 return SingleCellCoadd( 

327 outer=image_planes, 

328 psf=psf, 

329 inner_bbox=Box2I( 

330 corner=Point2I( 

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

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

333 ), 

334 dimensions=inner_cell_size, 

335 ), 

336 common=common, 

337 identifiers=identifiers, 

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

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

340 ) 

341 

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

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

344 

345 Returns 

346 ------- 

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

348 The WCS information read from the FITS file. 

349 """ 

350 # Read in WCS 

351 ps = PropertySet() 

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

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

354 wcs = afwGeom.makeSkyWcs(ps) 

355 return wcs 

356 

357 

358def writeMultipleCellCoaddAsFits( 

359 multiple_cell_coadd: MultipleCellCoadd, 

360 filename: str, 

361 overwrite: bool = False, 

362 metadata: PropertySet | None = None, 

363) -> None: 

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

365 

366 Parameters 

367 ---------- 

368 multiple_cell_coadd : `MultipleCellCoadd` 

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

370 filename : `str` 

371 The name of the file to write to. 

372 overwrite : `bool`, optional 

373 Whether to overwrite the file if it already exists? 

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

375 Additional metadata to write to the FITS file. 

376 

377 Notes 

378 ----- 

379 Changes to this function that modify the way the file is written to disk 

380 must be accompanied with a change to FILE_FORMAT_VERSION. 

381 """ 

382 cell_id = fits.Column( 

383 name="cell_id", 

384 format="2I", 

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

386 ) 

387 

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

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

390 image = fits.Column( 

391 name="image", 

392 unit=unit_array[0], 

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

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

395 array=image_array, 

396 ) 

397 

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

399 mask = fits.Column( 

400 name="mask", 

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

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

403 array=mask_array, 

404 ) 

405 

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

407 variance = fits.Column( 

408 name="variance", 

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

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

411 array=variance_array, 

412 ) 

413 

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

415 psf = fits.Column( 

416 name="psf", 

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

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

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

420 ) 

421 

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

423 hdu = fits.BinTableHDU.from_columns(col_defs) 

424 

425 grid_cell_size = multiple_cell_coadd.grid.cell_size 

426 grid_shape = multiple_cell_coadd.grid.shape 

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

428 grid_cards = { 

429 "GRCELL1": grid_cell_size.x, 

430 "GRCELL2": grid_cell_size.y, 

431 "GRSHAPE1": grid_shape.x, 

432 "GRSHAPE2": grid_shape.y, 

433 "GRMIN1": grid_min.x, 

434 "GRMIN2": grid_min.y, 

435 } 

436 hdu.header.extend(grid_cards) 

437 

438 outer_cell_size_cards = { 

439 "OCELL1": multiple_cell_coadd.outer_cell_size.x, 

440 "OCELL2": multiple_cell_coadd.outer_cell_size.y, 

441 } 

442 hdu.header.extend(outer_cell_size_cards) 

443 

444 psf_image_size_cards = { 

445 "PSFSIZE1": multiple_cell_coadd.psf_image_size.x, 

446 "PSFSIZE2": multiple_cell_coadd.psf_image_size.y, 

447 } 

448 hdu.header.extend(psf_image_size_cards) 

449 

450 inner_bbox_cards = { 

451 "INBBOX11": multiple_cell_coadd.inner_bbox.minX, 

452 "INBBOX12": multiple_cell_coadd.inner_bbox.minY, 

453 "INBBOX21": multiple_cell_coadd.inner_bbox.maxX, 

454 "INBBOX22": multiple_cell_coadd.inner_bbox.maxY, 

455 } 

456 hdu.header.extend(inner_bbox_cards) 

457 

458 wcs = multiple_cell_coadd.common.wcs 

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

460 primary_hdu = fits.PrimaryHDU() 

461 primary_hdu.header.extend(wcs_cards) 

462 

463 hdu.header["VERSION"] = FILE_FORMAT_VERSION 

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

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

466 # See DM-38843. 

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

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

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

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

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

472 

473 if metadata is not None: 

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

475 

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

477 hdu_list.writeto(filename, overwrite=overwrite)