Coverage for python / lsst / images / cells / _coadd.py: 38%
153 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-05-16 00:52 -0700
1# This file is part of lsst-images.
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# Use of this source code is governed by a 3-clause BSD-style
10# license that can be found in the LICENSE file.
12from __future__ import annotations
14__all__ = ("CellCoadd", "CellCoaddSerializationModel")
16import functools
17from collections.abc import Mapping, Sequence
18from types import EllipsisType
19from typing import TYPE_CHECKING, Any, cast
21import astropy.io.fits
22import astropy.units
23import astropy.wcs
24import pydantic
26from .._cell_grid import CellGrid, CellGridBounds, PatchDefinition
27from .._geom import YX, Box
28from .._image import Image, ImageSerializationModel
29from .._mask import Mask, MaskPlane, MaskSchema, MaskSerializationModel
30from .._masked_image import MaskedImage, MaskedImageSerializationModel
31from .._transforms import Projection, ProjectionSerializationModel, TractFrame
32from ..serialization import ArchiveReadError, InputArchive, OutputArchive
33from ._provenance import CoaddProvenance, CoaddProvenanceSerializationModel
34from ._psf import CellPointSpreadFunction, CellPointSpreadFunctionSerializationModel
36if TYPE_CHECKING:
37 try:
38 from lsst.cell_coadds import MultipleCellCoadd
39 from lsst.skymap import TractInfo
40 except ImportError:
41 type MultipleCellCoadd = Any # type: ignore[no-redef]
42 type TractInfo = Any # type: ignore[no-redef]
45class CellCoadd(MaskedImage):
46 """A coadd comprised of cells on a regular grid.
48 Parameters
49 ----------
50 image
51 The main image plane. If this has a `.Projection`, it will be used
52 for all planes unless a ``projection`` is passed separately.
53 mask
54 A bitmask image that annotates the main image plane. Must have the
55 same bounding box as ``image`` if provided. Any attached projection
56 is replaced (possibly by `None`).
57 variance
58 The per-pixel uncertainty of the main image as an image of variance
59 values. Must have the same bounding box as ``image`` if provided, and
60 its units must be the square of ``image.unit`` or `None`.
61 Values default to ``1.0``. Any attached projection is replaced
62 (possibly by `None`).
63 mask_fractions
64 A mapping from an input-image mask plane name to an image of the
65 weights sums of that plane.
66 noise_realizations
67 A sequence of images with Monte Carlo realizations of the noise in
68 the coadd.
69 mask_schema
70 Schema for the mask plane. Must be provided if and only if ``mask`` is
71 not provided.
72 projection
73 Projection that maps the pixel grid to the sky. Can only be `None` if
74 a projection is already attached to ``image``.
75 band
76 Name of the band.
77 psf
78 Effective point-spread function for the coadd. The missing cells
79 reported by ``psf.bounds`` are assumed to apply to all image data for
80 that cell as well (i.e. there is a PSF for a cell if and only if
81 there is image data for that cell).
82 patch
83 Identifiers and geometry of the full patch, if the image is confined
84 to a single patch. When present, the cell grid of the PSF and
85 provenance (if provideD) must be the full patch grid, even if its
86 bounds select a subset of that area.
87 provenance
88 Information about the images that went into the coadd.
89 """
91 def __init__(
92 self,
93 image: Image,
94 *,
95 mask: Mask | None = None,
96 variance: Image | None = None,
97 mask_fractions: Mapping[str, Image] | None = None,
98 noise_realizations: Sequence[Image] = (),
99 mask_schema: MaskSchema | None = None,
100 projection: Projection[TractFrame] | None = None,
101 band: str | None = None,
102 psf: CellPointSpreadFunction,
103 patch: PatchDefinition | None = None,
104 provenance: CoaddProvenance | None = None,
105 ):
106 super().__init__(
107 image,
108 mask=mask,
109 variance=variance,
110 mask_schema=mask_schema,
111 projection=projection,
112 )
113 if self.image.unit is None:
114 raise TypeError("The image component of a CellCoadd must have units.")
115 if self.image.projection is None:
116 raise TypeError("The projection component of a CellCoadd cannot be None.")
117 if not isinstance(self.image.projection.pixel_frame, TractFrame):
118 raise TypeError("The projection's pixel frame must be a TractFrame for CellCoadd.")
119 self._mask_fractions = dict(mask_fractions) if mask_fractions is not None else {}
120 self._noise_realizations = list(noise_realizations)
121 self._band = band
122 if psf.bounds.bbox != self.bbox:
123 psf = psf[self.bbox]
124 self._psf = psf
125 self._patch = patch
126 self._provenance = provenance
127 if self._provenance and not self._patch:
128 raise TypeError("A CellCoadd cannot carry provenance without a patch definition.")
130 @property
131 def skymap(self) -> str:
132 """Name of the skymap (`str`)."""
133 return self.projection.pixel_frame.skymap
135 @property
136 def tract(self) -> int:
137 """ID of the tract (`int`)."""
138 return self.projection.pixel_frame.tract
140 @property
141 def patch(self) -> PatchDefinition:
142 """Identifiers and geometry of the full patch, if the image is confined
143 to a single patch (`PatchDefinition`).
144 """
145 if self._patch is None:
146 raise AttributeError("Coadd has no patch information.")
147 return self._patch
149 @property
150 def band(self) -> str | None:
151 """Name of the band (`str` or `None`)."""
152 return self._band
154 @property
155 def mask_fractions(self) -> Mapping[str, Image]:
156 """A mapping from an input-image mask plane name to an image of the
157 weights sums of that plane
158 (`~collections.abc.Mapping` [`str`, `.Image`]).
159 """
160 return self._mask_fractions
162 @property
163 def noise_realizations(self) -> Sequence[Image]:
164 """A sequence of images with Monte Carlo realizations of the noise in
165 the coadd (`~collections.abc.Sequence` [`.Image`]).
166 """
167 return self._noise_realizations
169 @property
170 def unit(self) -> astropy.units.UnitBase:
171 """The units of the image plane (`astropy.units.Unit`)."""
172 return cast(astropy.units.UnitBase, super().unit)
174 @property
175 def projection(self) -> Projection[TractFrame]:
176 """The projection that maps the pixel grid to the sky
177 (`.Projection` [`.TractFrame`]).
178 """
179 return cast(Projection[TractFrame], super().projection)
181 @property
182 def psf(self) -> CellPointSpreadFunction:
183 """Effective point-spread function for the coadd
184 (`CellPointSpreadFunction`).
185 """
186 return self._psf
188 @property
189 def bounds(self) -> CellGridBounds:
190 """The grid of cells that overlap this coadd and a set of missing
191 cells (`CellGridBounds`).
192 """
193 return self._psf.bounds
195 @property
196 def grid(self) -> CellGrid:
197 """The grid of cells that overlap this coadd (`CellGrid`)."""
198 return self._psf.bounds.grid
200 @property
201 def provenance(self) -> CoaddProvenance:
202 """Information about the images that went into the coadd
203 (`CoaddProvenance` or `None`).
204 """
205 if self._provenance is None:
206 raise AttributeError("Coadd has no provenance information.")
207 return self._provenance
209 def __getitem__(self, bbox: Box | EllipsisType) -> CellCoadd:
210 if bbox is ...:
211 return self
212 super().__getitem__(bbox)
213 psf = self.psf[bbox]
214 return self._transfer_metadata(
215 CellCoadd(
216 self.image[bbox],
217 mask=self.mask[bbox],
218 variance=self.variance[bbox],
219 projection=self.projection,
220 mask_fractions={k: v[bbox] for k, v in self._mask_fractions.items()},
221 noise_realizations=[v[bbox] for v in self._noise_realizations],
222 band=self.band,
223 psf=psf,
224 patch=self._patch,
225 provenance=(
226 self._provenance.subset(psf.bounds.cell_indices())
227 if self._provenance is not None
228 else None
229 ),
230 ),
231 bbox=bbox,
232 )
234 def __str__(self) -> str:
235 return f"CellCoadd({self.bbox!s}, tract={self.tract})"
237 def __repr__(self) -> str:
238 return str(self)
240 def copy(self) -> CellCoadd:
241 """Deep-copy the coadd."""
242 return self._transfer_metadata(
243 CellCoadd(
244 image=self._image.copy(),
245 mask=self._mask.copy(),
246 variance=self._variance.copy(),
247 projection=self.projection,
248 mask_fractions={k: v.copy() for k, v in self._mask_fractions.items()},
249 noise_realizations=[v.copy() for v in self._noise_realizations],
250 band=self.band,
251 psf=self.psf,
252 patch=self.patch,
253 provenance=self.provenance,
254 ),
255 copy=True,
256 )
258 def serialize(self, archive: OutputArchive[Any]) -> CellCoaddSerializationModel:
259 """Serialize the image to an output archive.
261 Parameters
262 ----------
263 archive
264 Archive to write to.
265 """
266 serialized_image = archive.serialize_direct(
267 "image", functools.partial(self.image.serialize, save_projection=False)
268 )
269 serialized_mask = archive.serialize_direct(
270 "mask", functools.partial(self.mask.serialize, save_projection=False)
271 )
272 serialized_variance = archive.serialize_direct(
273 "variance", functools.partial(self.variance.serialize, save_projection=False)
274 )
275 serialized_projection = archive.serialize_direct("projection", self.projection.serialize)
276 serialized_mask_fractions = {
277 k: archive.serialize_direct(f"mask_fractions/{k}", v.serialize)
278 for k, v in self.mask_fractions.items()
279 }
280 serialized_noise_realizations = [
281 archive.serialize_direct(f"noise_realizations/{n}", v.serialize)
282 for n, v in enumerate(self.noise_realizations)
283 ]
284 serialized_psf = archive.serialize_direct("psf", self.psf.serialize)
285 serialized_provenance = (
286 archive.serialize_direct("provenance", self._provenance.serialize)
287 if self._provenance is not None
288 else None
289 )
290 return CellCoaddSerializationModel(
291 image=serialized_image,
292 mask=serialized_mask,
293 variance=serialized_variance,
294 projection=serialized_projection,
295 mask_fractions=serialized_mask_fractions,
296 noise_realizations=serialized_noise_realizations,
297 band=self._band,
298 psf=serialized_psf,
299 patch=self._patch,
300 provenance=serialized_provenance,
301 metadata=self.metadata,
302 )
304 @staticmethod
305 def _get_archive_tree_type[P: pydantic.BaseModel](
306 pointer_type: type[P],
307 ) -> type[CellCoaddSerializationModel[P]]:
308 """Return the serialization model type for this object for an archive
309 type that uses the given pointer type.
310 """
311 return CellCoaddSerializationModel[pointer_type] # type: ignore
313 # TODO: write_fits and read_fits inherited from MaskedImage, but that
314 # write_fits doesn't have compression-option kwargs for all of the new
315 # planes that CellCoadd adds. This makes me lean towards dropping the
316 # custom read_fits and write_fits in favor of the generic free functions
317 # in the fits subpackage, even though those aren't ideal either.
319 @staticmethod
320 def from_legacy( # type: ignore[override]
321 legacy: MultipleCellCoadd,
322 *,
323 plane_map: Mapping[str, MaskPlane] | None = None,
324 tract_info: TractInfo,
325 ) -> CellCoadd:
326 """Convert from an `lsst.cell_coadds.MultipleCellCoadd` instance.
328 Parameters
329 ----------
330 legacy
331 A `lsst.cell_coadds.MultipleCellCoadd` instance to convert.
332 plane_map
333 A mapping from legacy mask plane name to the new plane name and
334 description.
335 tract_info
336 Information about the full tract.
337 """
338 from lsst.geom import Box2I
340 legacy_bbox = Box2I()
341 for single_cell in legacy.cells.values():
342 legacy_bbox.include(single_cell.inner.bbox)
343 legacy_stitched = legacy.stitch(legacy_bbox)
344 unit = astropy.units.Unit(legacy.units.value)
345 tract_bbox = Box.from_legacy(tract_info.getBBox())
346 projection = Projection.from_legacy(
347 legacy.wcs,
348 TractFrame(
349 skymap=legacy.identifiers.skymap,
350 tract=legacy.identifiers.tract,
351 bbox=tract_bbox,
352 ),
353 pixel_bounds=tract_bbox,
354 )
355 band = legacy.identifiers.band
356 image = Image.from_legacy(legacy_stitched.image, unit=unit)
357 mask = Mask.from_legacy(legacy_stitched.mask, plane_map=plane_map)
358 variance = Image.from_legacy(legacy_stitched.variance, unit=unit**2)
359 noise_realizations = [
360 Image.from_legacy(noise_image) for noise_image in legacy_stitched.noise_realizations
361 ]
362 mask_fractions = (
363 {"rejected": Image.from_legacy(legacy_stitched.mask_fractions)}
364 if legacy_stitched.mask_fractions is not None
365 else {}
366 )
367 psf = CellPointSpreadFunction.from_legacy(legacy_stitched.psf, image.bbox)
368 patch_info = tract_info[legacy.identifiers.patch]
369 patch = PatchDefinition(
370 id=patch_info.getSequentialIndex(),
371 index=YX(y=legacy.identifiers.patch.y, x=legacy.identifiers.patch.x),
372 inner_bbox=Box.from_legacy(patch_info.getInnerBBox()),
373 cells=CellGrid.from_legacy(legacy.grid),
374 )
375 provenance = CoaddProvenance.from_legacy(legacy)
376 return CellCoadd(
377 image=image,
378 mask=mask,
379 variance=variance,
380 mask_fractions=mask_fractions,
381 noise_realizations=noise_realizations,
382 projection=projection,
383 band=band,
384 psf=psf,
385 patch=patch,
386 provenance=provenance,
387 )
390class CellCoaddSerializationModel[P: pydantic.BaseModel](MaskedImageSerializationModel[P]):
391 """A Pydantic model used to represent a serialized `CellCoadd`."""
393 # Inherited attributes are duplicated because that improves the docs
394 # (some limitation in the sphinx/pydantic integration), and these are
395 # important docs.
397 image: ImageSerializationModel[P] = pydantic.Field(description="The main data image.")
398 mask: MaskSerializationModel[P] = pydantic.Field(
399 description="Bitmask that annotates the main image's pixels."
400 )
401 variance: ImageSerializationModel[P] = pydantic.Field(
402 description="Per-pixel variance estimates for the main image."
403 )
404 projection: ProjectionSerializationModel[P] = pydantic.Field(
405 description="Projection that maps the pixel grid to the sky.",
406 )
407 mask_fractions: dict[str, ImageSerializationModel[P]] = pydantic.Field(
408 description=(
409 "A mapping from an input-image mask plane name to an image of the weights sums of that plane."
410 )
411 )
412 noise_realizations: list[ImageSerializationModel[P]] = pydantic.Field(
413 description=(
414 "A mapping from an input-image mask plane name to an image of the weights sums of that plane."
415 )
416 )
417 band: str | None = pydantic.Field(description="Name of the band.")
418 psf: CellPointSpreadFunctionSerializationModel = pydantic.Field(
419 description="Effective point-spread function model for the coadd."
420 )
421 patch: PatchDefinition | None = pydantic.Field(description="Identifiers and geometry for the patch.")
422 provenance: CoaddProvenanceSerializationModel | None = pydantic.Field(
423 description="Information about the images that went into the coadd."
424 )
426 def deserialize( # type: ignore[override]
427 self,
428 archive: InputArchive[Any],
429 *,
430 bbox: Box | None = None,
431 provenance: bool = True,
432 ) -> CellCoadd:
433 """Deserialize an image from an input archive.
435 Parameters
436 ----------
437 archive
438 Archive to read from.
439 bbox
440 Bounding box of a subimage to read instead.
441 provenance
442 Whether to read and attach provenance information.
443 """
444 masked_image = super().deserialize(archive, bbox=bbox)
445 mask_fractions = {
446 k.removeprefix("mask_fractions/"): v.deserialize(archive) for k, v in self.mask_fractions.items()
447 }
448 noise_realizations = [v.deserialize(archive) for v in self.noise_realizations]
449 projection = self.projection.deserialize(archive)
450 psf = self.psf.deserialize(archive, bbox=bbox)
451 coadd_provenance: CoaddProvenance | None = None
452 if self.provenance is not None and provenance:
453 coadd_provenance = self.provenance.deserialize(archive)
454 if bbox is not None:
455 coadd_provenance = coadd_provenance.subset(psf.bounds.cell_indices())
456 return CellCoadd(
457 masked_image.image,
458 mask=masked_image.mask,
459 variance=masked_image.variance,
460 mask_fractions=mask_fractions,
461 noise_realizations=noise_realizations,
462 projection=projection,
463 band=self.band,
464 psf=psf,
465 patch=self.patch,
466 provenance=coadd_provenance,
467 )._finish_deserialize(self)
469 def deserialize_psf(self, archive: InputArchive[Any], bbox: Box | None = None) -> CellPointSpreadFunction:
470 """Finish deserializing the PSF model."""
471 return self.psf.deserialize(archive, bbox=bbox)
473 def deserialize_provenance(self, archive: InputArchive[Any]) -> CoaddProvenance:
474 """Finish deserializing the provenance information."""
475 if self.provenance is not None:
476 return self.provenance.deserialize(archive)
477 raise ArchiveReadError("No coadd provenance stored in this file.")