Coverage for python / lsst / images / cells / _coadd.py: 37%

164 statements  

« prev     ^ index     » next       coverage.py v7.14.0, created at 2026-05-21 08:47 +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. 

11 

12from __future__ import annotations 

13 

14__all__ = ("CellCoadd", "CellCoaddSerializationModel") 

15 

16import functools 

17from collections.abc import Mapping, Sequence 

18from types import EllipsisType 

19from typing import TYPE_CHECKING, Any, cast 

20 

21import astropy.io.fits 

22import astropy.units 

23import astropy.wcs 

24import pydantic 

25 

26from .._backgrounds import BackgroundMap, BackgroundMapSerializationModel 

27from .._cell_grid import CellGrid, CellGridBounds, PatchDefinition 

28from .._geom import YX, Box 

29from .._image import Image, ImageSerializationModel 

30from .._mask import Mask, MaskPlane, MaskSchema, MaskSerializationModel 

31from .._masked_image import MaskedImage, MaskedImageSerializationModel 

32from .._transforms import Projection, ProjectionSerializationModel, TractFrame 

33from ..serialization import InputArchive, InvalidParameterError, OutputArchive 

34from ._provenance import CoaddProvenance, CoaddProvenanceSerializationModel 

35from ._psf import CellPointSpreadFunction, CellPointSpreadFunctionSerializationModel 

36 

37if TYPE_CHECKING: 

38 try: 

39 from lsst.cell_coadds import MultipleCellCoadd 

40 from lsst.skymap import TractInfo 

41 except ImportError: 

42 type MultipleCellCoadd = Any # type: ignore[no-redef] 

43 type TractInfo = Any # type: ignore[no-redef] 

44 

45 

46class CellCoadd(MaskedImage): 

47 """A coadd comprised of cells on a regular grid. 

48 

49 Parameters 

50 ---------- 

51 image 

52 The main image plane. If this has a `.Projection`, it will be used 

53 for all planes unless a ``projection`` is passed separately. 

54 mask 

55 A bitmask image that annotates the main image plane. Must have the 

56 same bounding box as ``image`` if provided. Any attached projection 

57 is replaced (possibly by `None`). 

58 variance 

59 The per-pixel uncertainty of the main image as an image of variance 

60 values. Must have the same bounding box as ``image`` if provided, and 

61 its units must be the square of ``image.unit`` or `None`. 

62 Values default to ``1.0``. Any attached projection is replaced 

63 (possibly by `None`). 

64 mask_fractions 

65 A mapping from an input-image mask plane name to an image of the 

66 weights sums of that plane. 

67 noise_realizations 

68 A sequence of images with Monte Carlo realizations of the noise in 

69 the coadd. 

70 mask_schema 

71 Schema for the mask plane. Must be provided if and only if ``mask`` is 

72 not provided. 

73 projection 

74 Projection that maps the pixel grid to the sky. Can only be `None` if 

75 a projection is already attached to ``image``. 

76 band 

77 Name of the band. 

78 psf 

79 Effective point-spread function for the coadd. The missing cells 

80 reported by ``psf.bounds`` are assumed to apply to all image data for 

81 that cell as well (i.e. there is a PSF for a cell if and only if 

82 there is image data for that cell). 

83 patch 

84 Identifiers and geometry of the full patch, if the image is confined 

85 to a single patch. When present, the cell grid of the PSF and 

86 provenance (if provideD) must be the full patch grid, even if its 

87 bounds select a subset of that area. 

88 provenance 

89 Information about the images that went into the coadd. 

90 backgrounds 

91 Background models associated with this image. 

92 """ 

93 

94 def __init__( 

95 self, 

96 image: Image, 

97 *, 

98 mask: Mask | None = None, 

99 variance: Image | None = None, 

100 mask_fractions: Mapping[str, Image] | None = None, 

101 noise_realizations: Sequence[Image] = (), 

102 mask_schema: MaskSchema | None = None, 

103 projection: Projection[TractFrame] | None = None, 

104 band: str | None = None, 

105 psf: CellPointSpreadFunction, 

106 patch: PatchDefinition | None = None, 

107 provenance: CoaddProvenance | None = None, 

108 backgrounds: BackgroundMap | None = None, 

109 ): 

110 super().__init__( 

111 image, 

112 mask=mask, 

113 variance=variance, 

114 mask_schema=mask_schema, 

115 projection=projection, 

116 ) 

117 if self.image.unit is None: 

118 raise TypeError("The image component of a CellCoadd must have units.") 

119 if self.image.projection is None: 

120 raise TypeError("The projection component of a CellCoadd cannot be None.") 

121 if not isinstance(self.image.projection.pixel_frame, TractFrame): 

122 raise TypeError("The projection's pixel frame must be a TractFrame for CellCoadd.") 

123 self._mask_fractions = dict(mask_fractions) if mask_fractions is not None else {} 

124 self._noise_realizations = list(noise_realizations) 

125 self._band = band 

126 if psf.bounds.bbox != self.bbox: 

127 psf = psf[self.bbox] 

128 self._psf = psf 

129 self._patch = patch 

130 self._provenance = provenance 

131 if self._provenance and not self._patch: 

132 raise TypeError("A CellCoadd cannot carry provenance without a patch definition.") 

133 self._backgrounds = backgrounds if backgrounds is not None else BackgroundMap() 

134 

135 @property 

136 def skymap(self) -> str: 

137 """Name of the skymap (`str`).""" 

138 return self.projection.pixel_frame.skymap 

139 

140 @property 

141 def tract(self) -> int: 

142 """ID of the tract (`int`).""" 

143 return self.projection.pixel_frame.tract 

144 

145 @property 

146 def patch(self) -> PatchDefinition: 

147 """Identifiers and geometry of the full patch, if the image is confined 

148 to a single patch (`PatchDefinition`). 

149 """ 

150 if self._patch is None: 

151 raise AttributeError("Coadd has no patch information.") 

152 return self._patch 

153 

154 @property 

155 def band(self) -> str | None: 

156 """Name of the band (`str` or `None`).""" 

157 return self._band 

158 

159 @property 

160 def mask_fractions(self) -> Mapping[str, Image]: 

161 """A mapping from an input-image mask plane name to an image of the 

162 weights sums of that plane 

163 (`~collections.abc.Mapping` [`str`, `.Image`]). 

164 """ 

165 return self._mask_fractions 

166 

167 @property 

168 def noise_realizations(self) -> Sequence[Image]: 

169 """A sequence of images with Monte Carlo realizations of the noise in 

170 the coadd (`~collections.abc.Sequence` [`.Image`]). 

171 """ 

172 return self._noise_realizations 

173 

174 @property 

175 def unit(self) -> astropy.units.UnitBase: 

176 """The units of the image plane (`astropy.units.Unit`).""" 

177 return cast(astropy.units.UnitBase, super().unit) 

178 

179 @property 

180 def projection(self) -> Projection[TractFrame]: 

181 """The projection that maps the pixel grid to the sky 

182 (`.Projection` [`.TractFrame`]). 

183 """ 

184 return cast(Projection[TractFrame], super().projection) 

185 

186 @property 

187 def psf(self) -> CellPointSpreadFunction: 

188 """Effective point-spread function for the coadd 

189 (`CellPointSpreadFunction`). 

190 """ 

191 return self._psf 

192 

193 @property 

194 def bounds(self) -> CellGridBounds: 

195 """The grid of cells that overlap this coadd and a set of missing 

196 cells (`CellGridBounds`). 

197 """ 

198 return self._psf.bounds 

199 

200 @property 

201 def grid(self) -> CellGrid: 

202 """The grid of cells that overlap this coadd (`CellGrid`).""" 

203 return self._psf.bounds.grid 

204 

205 @property 

206 def provenance(self) -> CoaddProvenance: 

207 """Information about the images that went into the coadd 

208 (`CoaddProvenance` or `None`). 

209 """ 

210 if self._provenance is None: 

211 raise AttributeError("Coadd has no provenance information.") 

212 return self._provenance 

213 

214 @property 

215 def backgrounds(self) -> BackgroundMap: 

216 """A mapping of backgrounds associated with this image 

217 (`.BackgroundMap`). 

218 """ 

219 return self._backgrounds 

220 

221 def __getitem__(self, bbox: Box | EllipsisType) -> CellCoadd: 

222 if bbox is ...: 

223 return self 

224 super().__getitem__(bbox) 

225 psf = self.psf[bbox] 

226 return self._transfer_metadata( 

227 CellCoadd( 

228 self.image[bbox], 

229 mask=self.mask[bbox], 

230 variance=self.variance[bbox], 

231 projection=self.projection, 

232 mask_fractions={k: v[bbox] for k, v in self._mask_fractions.items()}, 

233 noise_realizations=[v[bbox] for v in self._noise_realizations], 

234 band=self.band, 

235 psf=psf, 

236 patch=self._patch, 

237 provenance=( 

238 self._provenance.subset(psf.bounds.cell_indices()) 

239 if self._provenance is not None 

240 else None 

241 ), 

242 backgrounds=self._backgrounds, 

243 ), 

244 bbox=bbox, 

245 ) 

246 

247 def __str__(self) -> str: 

248 return f"CellCoadd({self.bbox!s}, tract={self.tract})" 

249 

250 def __repr__(self) -> str: 

251 return str(self) 

252 

253 def copy(self) -> CellCoadd: 

254 """Deep-copy the coadd.""" 

255 return self._transfer_metadata( 

256 CellCoadd( 

257 image=self._image.copy(), 

258 mask=self._mask.copy(), 

259 variance=self._variance.copy(), 

260 projection=self.projection, 

261 mask_fractions={k: v.copy() for k, v in self._mask_fractions.items()}, 

262 noise_realizations=[v.copy() for v in self._noise_realizations], 

263 band=self.band, 

264 psf=self.psf, 

265 patch=self.patch, 

266 provenance=self.provenance, 

267 backgrounds=self._backgrounds.copy(), 

268 ), 

269 copy=True, 

270 ) 

271 

272 def serialize(self, archive: OutputArchive[Any]) -> CellCoaddSerializationModel: 

273 """Serialize the image to an output archive. 

274 

275 Parameters 

276 ---------- 

277 archive 

278 Archive to write to. 

279 """ 

280 serialized_image = archive.serialize_direct( 

281 "image", functools.partial(self.image.serialize, save_projection=False) 

282 ) 

283 serialized_mask = archive.serialize_direct( 

284 "mask", functools.partial(self.mask.serialize, save_projection=False) 

285 ) 

286 serialized_variance = archive.serialize_direct( 

287 "variance", functools.partial(self.variance.serialize, save_projection=False) 

288 ) 

289 serialized_projection = archive.serialize_direct("projection", self.projection.serialize) 

290 serialized_mask_fractions = { 

291 k: archive.serialize_direct(f"mask_fractions/{k}", v.serialize) 

292 for k, v in self.mask_fractions.items() 

293 } 

294 serialized_noise_realizations = [ 

295 archive.serialize_direct(f"noise_realizations/{n}", v.serialize) 

296 for n, v in enumerate(self.noise_realizations) 

297 ] 

298 serialized_psf = archive.serialize_direct("psf", self.psf.serialize) 

299 serialized_provenance = ( 

300 archive.serialize_direct("provenance", self._provenance.serialize) 

301 if self._provenance is not None 

302 else None 

303 ) 

304 serialized_backgrounds = archive.serialize_direct("background", self._backgrounds.serialize) 

305 return CellCoaddSerializationModel( 

306 image=serialized_image, 

307 mask=serialized_mask, 

308 variance=serialized_variance, 

309 projection=serialized_projection, 

310 mask_fractions=serialized_mask_fractions, 

311 noise_realizations=serialized_noise_realizations, 

312 band=self._band, 

313 psf=serialized_psf, 

314 patch=self._patch, 

315 provenance=serialized_provenance, 

316 backgrounds=serialized_backgrounds, 

317 metadata=self.metadata, 

318 ) 

319 

320 @staticmethod 

321 def _get_archive_tree_type[P: pydantic.BaseModel]( 

322 pointer_type: type[P], 

323 ) -> type[CellCoaddSerializationModel[P]]: 

324 """Return the serialization model type for this object for an archive 

325 type that uses the given pointer type. 

326 """ 

327 return CellCoaddSerializationModel[pointer_type] # type: ignore 

328 

329 # TODO: write_fits and read_fits inherited from MaskedImage, but that 

330 # write_fits doesn't have compression-option kwargs for all of the new 

331 # planes that CellCoadd adds. This makes me lean towards dropping the 

332 # custom read_fits and write_fits in favor of the generic free functions 

333 # in the fits subpackage, even though those aren't ideal either. 

334 

335 @staticmethod 

336 def from_legacy( # type: ignore[override] 

337 legacy: MultipleCellCoadd, 

338 *, 

339 plane_map: Mapping[str, MaskPlane] | None = None, 

340 tract_info: TractInfo, 

341 ) -> CellCoadd: 

342 """Convert from an `lsst.cell_coadds.MultipleCellCoadd` instance. 

343 

344 Parameters 

345 ---------- 

346 legacy 

347 A `lsst.cell_coadds.MultipleCellCoadd` instance to convert. 

348 plane_map 

349 A mapping from legacy mask plane name to the new plane name and 

350 description. 

351 tract_info 

352 Information about the full tract. 

353 """ 

354 from lsst.geom import Box2I 

355 

356 legacy_bbox = Box2I() 

357 for single_cell in legacy.cells.values(): 

358 legacy_bbox.include(single_cell.inner.bbox) 

359 legacy_stitched = legacy.stitch(legacy_bbox) 

360 unit = astropy.units.Unit(legacy.units.value) 

361 tract_bbox = Box.from_legacy(tract_info.getBBox()) 

362 projection = Projection.from_legacy( 

363 legacy.wcs, 

364 TractFrame( 

365 skymap=legacy.identifiers.skymap, 

366 tract=legacy.identifiers.tract, 

367 bbox=tract_bbox, 

368 ), 

369 pixel_bounds=tract_bbox, 

370 ) 

371 band = legacy.identifiers.band 

372 image = Image.from_legacy(legacy_stitched.image, unit=unit) 

373 mask = Mask.from_legacy(legacy_stitched.mask, plane_map=plane_map) 

374 variance = Image.from_legacy(legacy_stitched.variance, unit=unit**2) 

375 noise_realizations = [ 

376 Image.from_legacy(noise_image) for noise_image in legacy_stitched.noise_realizations 

377 ] 

378 mask_fractions = ( 

379 {"rejected": Image.from_legacy(legacy_stitched.mask_fractions)} 

380 if legacy_stitched.mask_fractions is not None 

381 else {} 

382 ) 

383 psf = CellPointSpreadFunction.from_legacy(legacy_stitched.psf, image.bbox) 

384 patch_info = tract_info[legacy.identifiers.patch] 

385 patch = PatchDefinition( 

386 id=patch_info.getSequentialIndex(), 

387 index=YX(y=legacy.identifiers.patch.y, x=legacy.identifiers.patch.x), 

388 inner_bbox=Box.from_legacy(patch_info.getInnerBBox()), 

389 cells=CellGrid.from_legacy(legacy.grid), 

390 ) 

391 provenance = CoaddProvenance.from_legacy(legacy) 

392 return CellCoadd( 

393 image=image, 

394 mask=mask, 

395 variance=variance, 

396 mask_fractions=mask_fractions, 

397 noise_realizations=noise_realizations, 

398 projection=projection, 

399 band=band, 

400 psf=psf, 

401 patch=patch, 

402 provenance=provenance, 

403 ) 

404 

405 

406class CellCoaddSerializationModel[P: pydantic.BaseModel](MaskedImageSerializationModel[P]): 

407 """A Pydantic model used to represent a serialized `CellCoadd`.""" 

408 

409 # Inherited attributes are duplicated because that improves the docs 

410 # (some limitation in the sphinx/pydantic integration), and these are 

411 # important docs. 

412 

413 image: ImageSerializationModel[P] = pydantic.Field(description="The main data image.") 

414 mask: MaskSerializationModel[P] = pydantic.Field( 

415 description="Bitmask that annotates the main image's pixels." 

416 ) 

417 variance: ImageSerializationModel[P] = pydantic.Field( 

418 description="Per-pixel variance estimates for the main image." 

419 ) 

420 projection: ProjectionSerializationModel[P] = pydantic.Field( 

421 description="Projection that maps the pixel grid to the sky.", 

422 ) 

423 mask_fractions: dict[str, ImageSerializationModel[P]] = pydantic.Field( 

424 description=( 

425 "A mapping from an input-image mask plane name to an image of the weights sums of that plane." 

426 ) 

427 ) 

428 noise_realizations: list[ImageSerializationModel[P]] = pydantic.Field( 

429 description=( 

430 "A mapping from an input-image mask plane name to an image of the weights sums of that plane." 

431 ) 

432 ) 

433 band: str | None = pydantic.Field(description="Name of the band.") 

434 psf: CellPointSpreadFunctionSerializationModel = pydantic.Field( 

435 description="Effective point-spread function model for the coadd." 

436 ) 

437 patch: PatchDefinition | None = pydantic.Field(description="Identifiers and geometry for the patch.") 

438 provenance: CoaddProvenanceSerializationModel | None = pydantic.Field( 

439 description="Information about the images that went into the coadd." 

440 ) 

441 backgrounds: BackgroundMapSerializationModel = pydantic.Field( 

442 default_factory=BackgroundMapSerializationModel, 

443 description="Background models associated with this image.", 

444 ) 

445 

446 def deserialize( # type: ignore[override] 

447 self, 

448 archive: InputArchive[Any], 

449 *, 

450 bbox: Box | None = None, 

451 provenance: bool = True, 

452 **kwargs: Any, 

453 ) -> CellCoadd: 

454 """Deserialize an image from an input archive. 

455 

456 Parameters 

457 ---------- 

458 archive 

459 Archive to read from. 

460 bbox 

461 Bounding box of a subimage to read instead. 

462 provenance 

463 Whether to read and attach provenance information. 

464 **kwargs 

465 Unsupported keyword arguments are accepted only to provide better 

466 error messages (raising `.serialization.InvalidParameterError`). 

467 """ 

468 if kwargs: 

469 raise InvalidParameterError(f"Unrecognized parameters for CellCoadd: {set(kwargs.keys())}.") 

470 masked_image = super().deserialize(archive, bbox=bbox) 

471 mask_fractions = { 

472 k.removeprefix("mask_fractions/"): v.deserialize(archive) for k, v in self.mask_fractions.items() 

473 } 

474 noise_realizations = [v.deserialize(archive) for v in self.noise_realizations] 

475 projection = self.projection.deserialize(archive) 

476 psf = self.psf.deserialize(archive, bbox=bbox) 

477 coadd_provenance: CoaddProvenance | None = None 

478 if self.provenance is not None and provenance: 

479 coadd_provenance = self.provenance.deserialize(archive) 

480 if bbox is not None: 

481 coadd_provenance = coadd_provenance.subset(psf.bounds.cell_indices()) 

482 backgrounds = self.backgrounds.deserialize(archive) 

483 return CellCoadd( 

484 masked_image.image, 

485 mask=masked_image.mask, 

486 variance=masked_image.variance, 

487 mask_fractions=mask_fractions, 

488 noise_realizations=noise_realizations, 

489 projection=projection, 

490 band=self.band, 

491 psf=psf, 

492 patch=self.patch, 

493 provenance=coadd_provenance, 

494 backgrounds=backgrounds, 

495 )._finish_deserialize(self) 

496 

497 def deserialize_component(self, component: str, archive: InputArchive[Any], **kwargs: Any) -> Any: 

498 match component: 

499 case "mask_fractions": 

500 return { 

501 name: image_model.deserialize(archive, **kwargs) 

502 for name, image_model in self.mask_fractions.items() 

503 } 

504 case "noise_realizations": 

505 return [image_model.deserialize(archive, **kwargs) for image_model in self.noise_realizations] 

506 return super().deserialize_component(component, archive, **kwargs)