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

168 statements  

« prev     ^ index     » next       coverage.py v7.14.1, created at 2026-06-08 08:45 +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, ClassVar, 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 @staticmethod 

330 def from_legacy( # type: ignore[override] 

331 legacy: MultipleCellCoadd, 

332 *, 

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

334 tract_info: TractInfo, 

335 ) -> CellCoadd: 

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

337 

338 Parameters 

339 ---------- 

340 legacy 

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

342 plane_map 

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

344 description. 

345 tract_info 

346 Information about the full tract. 

347 """ 

348 from lsst.geom import Box2I 

349 

350 legacy_bbox = Box2I() 

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

352 legacy_bbox.include(single_cell.inner.bbox) 

353 legacy_stitched = legacy.stitch(legacy_bbox) 

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

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

356 projection = Projection.from_legacy( 

357 legacy.wcs, 

358 TractFrame( 

359 skymap=legacy.identifiers.skymap, 

360 tract=legacy.identifiers.tract, 

361 bbox=tract_bbox, 

362 ), 

363 pixel_bounds=tract_bbox, 

364 ) 

365 band = legacy.identifiers.band 

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

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

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

369 noise_realizations = [ 

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

371 ] 

372 mask_fractions = ( 

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

374 if legacy_stitched.mask_fractions is not None 

375 else {} 

376 ) 

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

378 patch_info = tract_info[legacy.identifiers.patch] 

379 patch = PatchDefinition( 

380 id=patch_info.getSequentialIndex(), 

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

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

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

384 ) 

385 provenance = CoaddProvenance.from_legacy(legacy) 

386 return CellCoadd( 

387 image=image, 

388 mask=mask, 

389 variance=variance, 

390 mask_fractions=mask_fractions, 

391 noise_realizations=noise_realizations, 

392 projection=projection, 

393 band=band, 

394 psf=psf, 

395 patch=patch, 

396 provenance=provenance, 

397 ) 

398 

399 

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

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

402 

403 SCHEMA_NAME: ClassVar[str] = "cell_coadd" 

404 SCHEMA_VERSION: ClassVar[str] = "1.0.0" 

405 MIN_READ_VERSION: ClassVar[int] = 1 

406 PUBLIC_TYPE: ClassVar[type] = CellCoadd 

407 

408 # Inherited attributes are duplicated because that improves the docs 

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

410 # important docs. 

411 

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

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

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

415 ) 

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

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

418 ) 

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

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

421 ) 

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

423 description=( 

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

425 ) 

426 ) 

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

428 description=( 

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

430 ) 

431 ) 

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

433 psf: CellPointSpreadFunctionSerializationModel = pydantic.Field( 

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

435 ) 

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

437 provenance: CoaddProvenanceSerializationModel | None = pydantic.Field( 

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

439 ) 

440 backgrounds: BackgroundMapSerializationModel = pydantic.Field( 

441 default_factory=BackgroundMapSerializationModel, 

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

443 ) 

444 

445 def deserialize( # type: ignore[override] 

446 self, 

447 archive: InputArchive[Any], 

448 *, 

449 bbox: Box | None = None, 

450 provenance: bool = True, 

451 **kwargs: Any, 

452 ) -> CellCoadd: 

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

454 

455 Parameters 

456 ---------- 

457 archive 

458 Archive to read from. 

459 bbox 

460 Bounding box of a subimage to read instead. 

461 provenance 

462 Whether to read and attach provenance information. 

463 **kwargs 

464 Unsupported keyword arguments are accepted only to provide better 

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

466 """ 

467 if kwargs: 

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

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

470 mask_fractions = { 

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

472 } 

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

474 projection = self.projection.deserialize(archive) 

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

476 coadd_provenance: CoaddProvenance | None = None 

477 if self.provenance is not None and provenance: 

478 coadd_provenance = self.provenance.deserialize(archive) 

479 if bbox is not None: 

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

481 backgrounds = self.backgrounds.deserialize(archive) 

482 return CellCoadd( 

483 masked_image.image, 

484 mask=masked_image.mask, 

485 variance=masked_image.variance, 

486 mask_fractions=mask_fractions, 

487 noise_realizations=noise_realizations, 

488 projection=projection, 

489 band=self.band, 

490 psf=psf, 

491 patch=self.patch, 

492 provenance=coadd_provenance, 

493 backgrounds=backgrounds, 

494 )._finish_deserialize(self) 

495 

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

497 match component: 

498 case "mask_fractions": 

499 return { 

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

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

502 } 

503 case "noise_realizations": 

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

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