Coverage for python/lsst/cell_coadds/_fits.py: 24%
95 statements
« prev ^ index » next coverage.py v7.4.1, created at 2024-02-09 12:31 +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/>.
22from __future__ import annotations
24__all__ = (
25 "CellCoaddFitsFormatter",
26 "CellCoaddFitsReader",
27 "writeMultipleCellCoaddAsFits",
28)
30import os
31from collections.abc import Mapping
32from typing import Any
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
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
51class CellCoaddFitsFormatter(FitsGenericFormatter):
52 """Interface for writing and reading cell coadds to/from FITS files.
54 This assumes the existence of readFits and writeFits methods (for now).
55 """
58class CellCoaddFitsReader:
59 """A reader class to read from a FITS file and produce cell-based coadds.
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).
66 Parameters
67 ----------
68 filename : `str`
69 The name of the FITS file to read.
70 """
72 def __init__(self, filename: str) -> None:
73 if not os.path.exists(filename):
74 raise FileNotFoundError(f"File {filename} not found")
76 self.filename = filename
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
84 # Read in WCS
85 ps = PropertySet()
86 ps.update(hdu_list[0].header)
87 wcs = afwGeom.makeSkyWcs(ps)
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 )
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)
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 )
113 outer_cell_size = Extent2I(header["OCELL1"], header["OCELL2"])
114 psf_image_size = Extent2I(header["PSFSIZE1"], header["PSFSIZE2"])
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 )
135 return coadd
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.
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.
163 Returns
164 -------
165 coadd : `SingleCellCoadd`
166 The coadd read from the file.
167 """
168 buffer = (outer_cell_size - inner_cell_size) // 2
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 )
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 )
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 )
214 def readWcs(self) -> afwGeom.SkyWcs:
215 """Read the WCS information from the FITS file.
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
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.
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 )
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 )
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 )
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 )
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 )
289 col_defs = fits.ColDefs([cell_id, image, mask, variance, psf])
290 hdu = fits.BinTableHDU.from_columns(col_defs)
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)
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)
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)
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)
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)
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
339 if metadata is not None:
340 hdu.header.extend(metadata.toDict())
342 hdu_list = fits.HDUList([primary_hdu, hdu])
343 hdu_list.writeto(filename, overwrite=overwrite)