Coverage for python / lsst / images / cells / _coadd.py: 39%
154 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-05-07 08:34 +0000
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 # Type-checkers want the model argument to only require
305 # MaskedImageSerializationModel[Any], and they'd be absolutely right if
306 # this were a regular instance method. But whether Liskov substitution
307 # applies to classmethods and staticmethods is sort of context-dependent,
308 # and here we do not want it to.
309 @staticmethod
310 def deserialize( # type: ignore[override]
311 model: CellCoaddSerializationModel[Any],
312 archive: InputArchive[Any],
313 *,
314 bbox: Box | None = None,
315 provenance: bool = True,
316 ) -> CellCoadd:
317 """Deserialize an image from an input archive.
319 Parameters
320 ----------
321 model
322 A Pydantic model representation of the image, holding references
323 to data stored in the archive.
324 archive
325 Archive to read from.
326 bbox
327 Bounding box of a subimage to read instead.
328 provenance
329 Whether to read and attach provenance information.
330 """
331 masked_image = MaskedImage.deserialize(model, archive, bbox=bbox)
332 mask_fractions = {
333 k.removeprefix("mask_fractions/"): Image.deserialize(v, archive)
334 for k, v in model.mask_fractions.items()
335 }
336 noise_realizations = [Image.deserialize(v, archive) for v in model.noise_realizations]
337 projection = Projection.deserialize(model.projection, archive)
338 psf = CellPointSpreadFunction.deserialize(model.psf, archive, bbox=bbox)
339 coadd_provenance: CoaddProvenance | None = None
340 if model.provenance is not None and provenance:
341 coadd_provenance = CoaddProvenance.deserialize(model.provenance, archive)
342 if bbox is not None:
343 coadd_provenance = coadd_provenance.subset(psf.bounds.cell_indices())
344 return CellCoadd(
345 masked_image.image,
346 mask=masked_image.mask,
347 variance=masked_image.variance,
348 mask_fractions=mask_fractions,
349 noise_realizations=noise_realizations,
350 projection=projection,
351 band=model.band,
352 psf=psf,
353 patch=model.patch,
354 provenance=coadd_provenance,
355 )._finish_deserialize(model)
357 @staticmethod
358 def _get_archive_tree_type[P: pydantic.BaseModel](
359 pointer_type: type[P],
360 ) -> type[CellCoaddSerializationModel[P]]:
361 """Return the serialization model type for this object for an archive
362 type that uses the given pointer type.
363 """
364 return CellCoaddSerializationModel[pointer_type] # type: ignore
366 # TODO: write_fits and read_fits inherited from MaskedImage, but that
367 # write_fits doesn't have compression-option kwargs for all of the new
368 # planes that CellCoadd adds. This makes me lean towards dropping the
369 # custom read_fits and write_fits in favor of the generic free functions
370 # in the fits subpackage, even though those aren't ideal either.
372 @staticmethod
373 def from_legacy( # type: ignore[override]
374 legacy: MultipleCellCoadd,
375 *,
376 plane_map: Mapping[str, MaskPlane] | None = None,
377 tract_info: TractInfo,
378 ) -> CellCoadd:
379 """Convert from an `lsst.cell_coadds.MultipleCellCoadd` instance.
381 Parameters
382 ----------
383 legacy
384 A `lsst.cell_coadds.MultipleCellCoadd` instance to convert.
385 plane_map
386 A mapping from legacy mask plane name to the new plane name and
387 description.
388 tract_info
389 Information about the full tract.
390 """
391 from lsst.geom import Box2I
393 legacy_bbox = Box2I()
394 for single_cell in legacy.cells.values():
395 legacy_bbox.include(single_cell.inner.bbox)
396 legacy_stitched = legacy.stitch(legacy_bbox)
397 unit = astropy.units.Unit(legacy.units.value)
398 tract_bbox = Box.from_legacy(tract_info.getBBox())
399 projection = Projection.from_legacy(
400 legacy.wcs,
401 TractFrame(
402 skymap=legacy.identifiers.skymap,
403 tract=legacy.identifiers.tract,
404 bbox=tract_bbox,
405 ),
406 pixel_bounds=tract_bbox,
407 )
408 band = legacy.identifiers.band
409 image = Image.from_legacy(legacy_stitched.image, unit=unit)
410 mask = Mask.from_legacy(legacy_stitched.mask, plane_map=plane_map)
411 variance = Image.from_legacy(legacy_stitched.variance, unit=unit**2)
412 noise_realizations = [
413 Image.from_legacy(noise_image) for noise_image in legacy_stitched.noise_realizations
414 ]
415 mask_fractions = (
416 {"rejected": Image.from_legacy(legacy_stitched.mask_fractions)}
417 if legacy_stitched.mask_fractions is not None
418 else {}
419 )
420 psf = CellPointSpreadFunction.from_legacy(legacy_stitched.psf, image.bbox)
421 patch_info = tract_info[legacy.identifiers.patch]
422 patch = PatchDefinition(
423 id=patch_info.getSequentialIndex(),
424 index=YX(y=legacy.identifiers.patch.y, x=legacy.identifiers.patch.x),
425 inner_bbox=Box.from_legacy(patch_info.getInnerBBox()),
426 cells=CellGrid.from_legacy(legacy.grid),
427 )
428 provenance = CoaddProvenance.from_legacy(legacy)
429 return CellCoadd(
430 image=image,
431 mask=mask,
432 variance=variance,
433 mask_fractions=mask_fractions,
434 noise_realizations=noise_realizations,
435 projection=projection,
436 band=band,
437 psf=psf,
438 patch=patch,
439 provenance=provenance,
440 )
443class CellCoaddSerializationModel[P: pydantic.BaseModel](MaskedImageSerializationModel[P]):
444 """A Pydantic model used to represent a serialized `CellCoadd`."""
446 # Inherited attributes are duplicated because that improves the docs
447 # (some limitation in the sphinx/pydantic integration), and these are
448 # important docs.
450 image: ImageSerializationModel[P] = pydantic.Field(description="The main data image.")
451 mask: MaskSerializationModel[P] = pydantic.Field(
452 description="Bitmask that annotates the main image's pixels."
453 )
454 variance: ImageSerializationModel[P] = pydantic.Field(
455 description="Per-pixel variance estimates for the main image."
456 )
457 projection: ProjectionSerializationModel[P] = pydantic.Field(
458 description="Projection that maps the pixel grid to the sky.",
459 )
460 mask_fractions: dict[str, ImageSerializationModel[P]] = pydantic.Field(
461 description=(
462 "A mapping from an input-image mask plane name to an image of the weights sums of that plane."
463 )
464 )
465 noise_realizations: list[ImageSerializationModel[P]] = pydantic.Field(
466 description=(
467 "A mapping from an input-image mask plane name to an image of the weights sums of that plane."
468 )
469 )
470 band: str | None = pydantic.Field(description="Name of the band.")
471 psf: CellPointSpreadFunctionSerializationModel = pydantic.Field(
472 description="Effective point-spread function model for the coadd."
473 )
474 patch: PatchDefinition | None = pydantic.Field(description="Identifiers and geometry for the patch.")
475 provenance: CoaddProvenanceSerializationModel | None = pydantic.Field(
476 description="Information about the images that went into the coadd."
477 )
479 def deserialize_psf(self, archive: InputArchive[Any], bbox: Box | None = None) -> CellPointSpreadFunction:
480 """Finish deserializing the PSF model."""
481 return CellPointSpreadFunction.deserialize(self.psf, archive, bbox=bbox)
483 def deserialize_provenance(self, archive: InputArchive[Any]) -> CoaddProvenance:
484 """Finish deserializing the provenance information."""
485 if self.provenance is not None:
486 return CoaddProvenance.deserialize(self.provenance, archive)
487 raise ArchiveReadError("No coadd provenance stored in this file.")