Coverage for python / lsst / images / _visit_image.py: 28%

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__ = ("VisitImage", "VisitImageSerializationModel") 

15 

16import functools 

17import warnings 

18from collections.abc import Callable, Mapping, MutableMapping 

19from types import EllipsisType 

20from typing import Any, Literal, cast, overload 

21 

22import astropy.io.fits 

23import astropy.units 

24import astropy.wcs 

25import pydantic 

26from astro_metadata_translator import ObservationInfo, VisitInfoTranslator 

27 

28from ._geom import Box 

29from ._image import Image, ImageSerializationModel 

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

31from ._masked_image import MaskedImage, MaskedImageSerializationModel 

32from ._transforms import DetectorFrame, Projection, ProjectionAstropyView, ProjectionSerializationModel 

33from .fits import FitsOpaqueMetadata 

34from .psfs import ( 

35 PiffSerializationModel, 

36 PiffWrapper, 

37 PointSpreadFunction, 

38 PSFExSerializationModel, 

39 PSFExWrapper, 

40) 

41from .serialization import ArchiveReadError, InputArchive, MetadataValue, OutputArchive 

42 

43 

44def _obs_info_from_md(md: MutableMapping[str, Any], visit_info: Any = None) -> ObservationInfo: 

45 # Try to get an ObservationInfo from the primary header as if 

46 # it's a raw header. Else fallback. 

47 try: 

48 obs_info = ObservationInfo.from_header(md, quiet=True) 

49 except ValueError: 

50 # Not known translator. Must fall back to visit info. If we have 

51 # an actual VisitInfo, serialize it since we know that it will be 

52 # complete. 

53 if visit_info is not None: 

54 from lsst.afw.image import setVisitInfoMetadata 

55 from lsst.daf.base import PropertyList 

56 

57 pl = PropertyList() 

58 setVisitInfoMetadata(pl, visit_info) 

59 # Merge so that we still have access to butler provenance. 

60 md.update(pl) 

61 

62 # Try the given header looking for VisitInfo hints. 

63 # We get lots of warnings if nothing can be found. Currently 

64 # no way to disable those without capturing them. 

65 obs_info = ObservationInfo.from_header(md, translator_class=VisitInfoTranslator, quiet=True) 

66 return obs_info 

67 

68 

69class VisitImage(MaskedImage): 

70 """A calibrated single-visit image. 

71 

72 Parameters 

73 ---------- 

74 image 

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

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

77 mask 

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

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

80 is replaced (possibly by `None`). 

81 variance 

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

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

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

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

86 (possibly by `None`). 

87 mask_schema 

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

89 not provided. 

90 opaque_metadata 

91 Opaque metadata obtained from reading this object from storage. It may 

92 be provided when writing to storage to propagate that metadata and/or 

93 preserve file-format-specific options (e.g. compression parameters). 

94 projection 

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

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

97 psf 

98 Point-spread function model for this image, or an exception explaining 

99 why it could not be read (to be raised if the PSF is requested later). 

100 obs_info 

101 General information about this visit in standardized form. 

102 metadata 

103 Arbitrary flexible metadata to associate with the image. 

104 """ 

105 

106 def __init__( 

107 self, 

108 image: Image, 

109 *, 

110 mask: Mask | None = None, 

111 variance: Image | None = None, 

112 mask_schema: MaskSchema | None = None, 

113 projection: Projection[DetectorFrame] | None = None, 

114 obs_info: ObservationInfo | None = None, 

115 psf: PointSpreadFunction | ArchiveReadError, 

116 metadata: dict[str, MetadataValue] | None = None, 

117 ): 

118 super().__init__( 

119 image, 

120 mask=mask, 

121 variance=variance, 

122 mask_schema=mask_schema, 

123 projection=projection, 

124 obs_info=obs_info, 

125 metadata=metadata, 

126 ) 

127 if self.image.unit is None: 

128 raise TypeError("The image component of a VisitImage must have units.") 

129 if self.image.projection is None: 

130 raise TypeError("The projection component of a VisitImage cannot be None.") 

131 if self.image.obs_info is None: 

132 raise TypeError("The observation info component of a VisitImage cannot be None.") 

133 if not isinstance(self.image.projection.pixel_frame, DetectorFrame): 

134 raise TypeError("The projection's pixel frame must be a DetectorFrame for VisitImage.") 

135 self._psf = psf 

136 

137 @property 

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

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

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

141 

142 @property 

143 def projection(self) -> Projection[DetectorFrame]: 

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

145 (`Projection` [`DetectorFrame`]). 

146 """ 

147 return cast(Projection[DetectorFrame], super().projection) 

148 

149 @property 

150 def obs_info(self) -> ObservationInfo: 

151 """General information about this observation in standard form. 

152 (`~astro_metadata_translator.ObservationInfo`). 

153 """ 

154 obs_info = self.image.obs_info 

155 assert obs_info is not None 

156 return obs_info 

157 

158 @property 

159 def astropy_wcs(self) -> ProjectionAstropyView: 

160 """An Astropy WCS for the pixel arrays (`ProjectionAstropyView`). 

161 

162 Notes 

163 ----- 

164 As expected for Astropy WCS objects, this defines pixel coordinates 

165 such that the first row and column in the arrays are ``(0, 0)``, not 

166 ``bbox.start``, as is the case for `projection`. 

167 

168 This object satisfies the `astropy.wcs.wcsapi.BaseHighLevelWCS` and 

169 `astropy.wcs.wcsapi.BaseLowLevelWCS` interfaces, but it is not an 

170 `astropy.wcs.WCS` (use `fits_wcs` for that). 

171 """ 

172 return cast(ProjectionAstropyView, super().astropy_wcs) 

173 

174 @property 

175 def psf(self) -> PointSpreadFunction: 

176 """The point-spread function model for this image 

177 (`.psfs.PointSpreadFunction`). 

178 """ 

179 if isinstance(self._psf, ArchiveReadError): 

180 raise self._psf 

181 return self._psf 

182 

183 def __getitem__(self, bbox: Box | EllipsisType) -> VisitImage: 

184 super().__getitem__(bbox) 

185 if bbox is ...: 

186 return self 

187 return self._transfer_metadata( 

188 VisitImage( 

189 self.image[bbox], 

190 mask=self.mask[bbox], 

191 variance=self.variance[bbox], 

192 projection=self.projection, 

193 psf=self.psf, 

194 obs_info=self.obs_info, 

195 ), 

196 bbox=bbox, 

197 ) 

198 

199 def __str__(self) -> str: 

200 return f"VisitImage({self.image!s}, {list(self.mask.schema.names)})" 

201 

202 def __repr__(self) -> str: 

203 return f"VisitImage({self.image!r}, mask_schema={self.mask.schema!r})" 

204 

205 def copy(self) -> VisitImage: 

206 """Deep-copy the visit image.""" 

207 return self._transfer_metadata( 

208 VisitImage( 

209 image=self._image.copy(), 

210 mask=self._mask.copy(), 

211 variance=self._variance.copy(), 

212 psf=self._psf, 

213 obs_info=self.obs_info, 

214 ), 

215 copy=True, 

216 ) 

217 

218 def serialize(self, archive: OutputArchive[Any]) -> VisitImageSerializationModel: 

219 serialized_image = archive.serialize_direct( 

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

221 ) 

222 serialized_mask = archive.serialize_direct( 

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

224 ) 

225 serialized_variance = archive.serialize_direct( 

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

227 ) 

228 serialized_projection = ( 

229 archive.serialize_direct("projection", self.projection.serialize) 

230 if self.projection is not None 

231 else None 

232 ) 

233 serialized_psf: PiffSerializationModel | PSFExSerializationModel 

234 match self._psf: 

235 # MyPy is able to figure things out here with this match statement, 

236 # but not a single isinstance check on both types. 

237 case PiffWrapper(): 

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

239 case PSFExWrapper(): 

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

241 case _: 

242 raise TypeError( 

243 f"Cannot serialize VisitImage with unrecognized PSF type {type(self._psf).__name__}." 

244 ) 

245 return VisitImageSerializationModel( 

246 image=serialized_image, 

247 mask=serialized_mask, 

248 variance=serialized_variance, 

249 projection=serialized_projection, 

250 psf=serialized_psf, 

251 obs_info=self.obs_info, 

252 metadata=self.metadata, 

253 ) 

254 

255 # Type-checkers want the model argument to only require 

256 # MaskedImageSerializationModel[Any], and they'd be absolutely right if 

257 # this were a regular instance method. But whether Liskov substitution 

258 # applies to classmethods and staticmethods is sort of context-dependent, 

259 # and here we do not want it to. 

260 @staticmethod 

261 def deserialize( 

262 model: VisitImageSerializationModel[Any], # type: ignore[override] 

263 archive: InputArchive[Any], 

264 *, 

265 bbox: Box | None = None, 

266 ) -> VisitImage: 

267 masked_image = MaskedImage.deserialize(model, archive, bbox=bbox) 

268 psf = model.deserialize_psf(archive) 

269 projection = Projection.deserialize(model.projection, archive) 

270 return VisitImage( 

271 masked_image.image, 

272 mask=masked_image.mask, 

273 variance=masked_image.variance, 

274 psf=psf, 

275 projection=projection, 

276 obs_info=model.obs_info, 

277 )._finish_deserialize(model) 

278 

279 @staticmethod 

280 def _get_archive_tree_type[P: pydantic.BaseModel]( 

281 pointer_type: type[P], 

282 ) -> type[VisitImageSerializationModel[P]]: 

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

284 type that uses the given pointer type. 

285 """ 

286 return VisitImageSerializationModel[pointer_type] # type: ignore 

287 

288 # write_fits and read_fits inherited from MaskedImage. 

289 

290 @staticmethod 

291 def from_legacy( 

292 legacy: Any, 

293 *, 

294 unit: astropy.units.Unit | None = None, 

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

296 instrument: str | None = None, 

297 visit: int | None = None, 

298 component: str | None = None, 

299 ) -> VisitImage: 

300 """Convert from an `lsst.afw.image.Exposure` instance. 

301 

302 Parameters 

303 ---------- 

304 legacy 

305 An `lsst.afw.image.Exposure` instance that will share image and 

306 variance (but not mask) pixel data with the returned object. 

307 unit 

308 Units of the image. If not provided, the ``BUNIT`` metadata 

309 key will be used, if available. 

310 plane_map 

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

312 description. If `None` (default) 

313 `get_legacy_visit_image_mask_planes` is used. 

314 instrument 

315 Name of the instrument. Extracted from the metadata if not 

316 provided. 

317 visit 

318 ID of the visit. Extracted from the metadata if not provided. 

319 """ 

320 if plane_map is None: 

321 plane_map = get_legacy_visit_image_mask_planes() 

322 md = legacy.getMetadata() 

323 obs_info = _obs_info_from_md(md, visit_info=legacy.info.getVisitInfo()) 

324 instrument = _extract_or_check_header("LSST BUTLER DATAID INSTRUMENT", instrument, md, str) 

325 visit = _extract_or_check_header("LSST BUTLER DATAID VISIT", visit, md, int) 

326 unit = _extract_or_check_header("BUNIT", unit, md, lambda x: astropy.units.Unit(x, format="fits")) 

327 legacy_wcs = legacy.getWcs() 

328 if legacy_wcs is None: 

329 raise ValueError("Exposure does not have a SkyWcs.") 

330 legacy_detector = legacy.getDetector() 

331 if legacy_detector is None: 

332 raise ValueError("Exposure does not have a Detector.") 

333 detector_bbox = Box.from_legacy(legacy_detector.getBBox()) 

334 opaque_fits_metadata = FitsOpaqueMetadata() 

335 primary_header = astropy.io.fits.Header() 

336 with warnings.catch_warnings(): 

337 # Silence warnings about long keys becoming HIERARCH. 

338 warnings.simplefilter("ignore", category=astropy.io.fits.verify.VerifyWarning) 

339 primary_header.update(md.toOrderedDict()) 

340 opaque_fits_metadata.extract_legacy_primary_header(primary_header) 

341 projection = Projection.from_legacy( 

342 legacy_wcs, 

343 DetectorFrame( 

344 instrument=instrument, 

345 visit=visit, 

346 detector=legacy_detector.getId(), 

347 bbox=detector_bbox, 

348 ), 

349 ) 

350 legacy_psf = legacy.getPsf() 

351 if legacy_psf is None: 

352 raise ValueError("Exposure file does not have a Psf.") 

353 psf = PointSpreadFunction.from_legacy(legacy_psf, bounds=detector_bbox) 

354 masked_image = MaskedImage.from_legacy(legacy.getMaskedImage(), unit=unit, plane_map=plane_map) 

355 result = VisitImage( 

356 image=masked_image.image.view(unit=unit), 

357 mask=masked_image.mask, 

358 variance=masked_image.variance, 

359 projection=projection, 

360 psf=psf, 

361 obs_info=obs_info, 

362 ) 

363 result._opaque_metadata = masked_image._opaque_metadata 

364 return result 

365 

366 @overload # type: ignore[override] 

367 @staticmethod 

368 def read_legacy( 368 ↛ exitline 368 didn't return from function 'read_legacy' because

369 filename: str, 

370 *, 

371 component: Literal["bbox"], 

372 ) -> Box: ... 

373 

374 @overload 

375 @staticmethod 

376 def read_legacy( 376 ↛ exitline 376 didn't return from function 'read_legacy' because

377 filename: str, 

378 *, 

379 preserve_quantization: bool = False, 

380 instrument: str | None = None, 

381 visit: int | None = None, 

382 component: Literal["image"], 

383 ) -> Image: ... 

384 

385 @overload 

386 @staticmethod 

387 def read_legacy( 387 ↛ exitline 387 didn't return from function 'read_legacy' because

388 filename: str, 

389 *, 

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

391 instrument: str | None = None, 

392 visit: int | None = None, 

393 component: Literal["mask"], 

394 ) -> Mask: ... 

395 

396 @overload 

397 @staticmethod 

398 def read_legacy( 398 ↛ exitline 398 didn't return from function 'read_legacy' because

399 filename: str, 

400 *, 

401 preserve_quantization: bool = False, 

402 instrument: str | None = None, 

403 visit: int | None = None, 

404 component: Literal["variance"], 

405 ) -> Image: ... 

406 

407 @overload 

408 @staticmethod 

409 def read_legacy( 409 ↛ exitline 409 didn't return from function 'read_legacy' because

410 filename: str, 

411 *, 

412 instrument: str | None = None, 

413 visit: int | None = None, 

414 component: Literal["projection"], 

415 ) -> Projection[DetectorFrame]: ... 

416 

417 @overload 

418 @staticmethod 

419 def read_legacy( 419 ↛ exitline 419 didn't return from function 'read_legacy' because

420 filename: str, 

421 *, 

422 component: Literal["psf"], 

423 ) -> PointSpreadFunction: ... 

424 

425 @overload 

426 @staticmethod 

427 def read_legacy( 427 ↛ exitline 427 didn't return from function 'read_legacy' because

428 filename: str, 

429 *, 

430 preserve_quantization: bool = False, 

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

432 instrument: str | None = None, 

433 visit: int | None = None, 

434 component: None = None, 

435 ) -> VisitImage: ... 

436 

437 @staticmethod 

438 def read_legacy( # type: ignore[override] 

439 filename: str, 

440 *, 

441 preserve_quantization: bool = False, 

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

443 instrument: str | None = None, 

444 visit: int | None = None, 

445 component: Literal["bbox", "image", "mask", "variance", "projection", "psf"] | None = None, 

446 ) -> Any: 

447 """Read a FITS file written by `lsst.afw.image.Exposure.writeFits`. 

448 

449 Parameters 

450 ---------- 

451 filename 

452 Full name of the file. 

453 preserve_quantization 

454 If `True`, ensure that writing the masked image back out again will 

455 exactly preserve quantization-compressed pixel values. This causes 

456 the image and variance plane arrays to be marked as read-only and 

457 stores the original binary table data for those planes in memory. 

458 If the `MaskedImage` is copied, the precompressed pixel values are 

459 not transferred to the copy. 

460 plane_map 

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

462 description. If `None` (default) 

463 `get_legacy_visit_image_mask_planes` is used. 

464 instrument 

465 Name of the instrument. Read from the primary header if not 

466 provided. 

467 visit 

468 ID of the visit. Read from the primary header if not 

469 provided. 

470 component 

471 A component to read instead of the full image. 

472 """ 

473 from lsst.afw.image import ExposureFitsReader 

474 

475 reader = ExposureFitsReader(filename) 

476 if component == "bbox": 

477 return Box.from_legacy(reader.readBBox()) 

478 legacy_detector = reader.readDetector() 

479 if legacy_detector is None: 

480 raise ValueError(f"Exposure file {filename!r} does not have a Detector.") 

481 detector_bbox = Box.from_legacy(legacy_detector.getBBox()) 

482 if component in (None, "image", "mask", "variance", "projection"): 

483 legacy_wcs = reader.readWcs() 

484 if legacy_wcs is None: 

485 raise ValueError(f"Exposure file {filename!r} does not have a SkyWcs.") 

486 if component in (None, "psf"): 

487 legacy_psf = reader.readPsf() 

488 if legacy_psf is None: 

489 raise ValueError(f"Exposure file {filename!r} does not have a Psf.") 

490 psf = PointSpreadFunction.from_legacy(legacy_psf, bounds=detector_bbox) 

491 if component == "psf": 

492 return psf 

493 assert component in (None, "image", "mask", "variance", "projection"), component # for MyPy 

494 with astropy.io.fits.open(filename) as hdu_list: 

495 primary_header = hdu_list[0].header 

496 if instrument is None: 

497 try: 

498 instrument = primary_header["LSST BUTLER DATAID INSTRUMENT"] 

499 except LookupError: 

500 raise ValueError( 

501 "Instrument could not be found in butler data ID FITS headers and must be provided." 

502 ) from None 

503 if visit is None: 

504 try: 

505 visit = primary_header["LSST BUTLER DATAID VISIT"] 

506 except LookupError: 

507 raise ValueError( 

508 "Visit ID could not be found in butler data ID FITS headers and must be provided." 

509 ) from None 

510 projection = Projection.from_legacy( 

511 legacy_wcs, 

512 DetectorFrame( 

513 instrument=instrument, 

514 visit=visit, 

515 detector=legacy_detector.getId(), 

516 bbox=detector_bbox, 

517 ), 

518 ) 

519 if component == "projection": 

520 return projection 

521 if plane_map is None: 

522 plane_map = get_legacy_visit_image_mask_planes() 

523 assert component != "psf", component # for MyPy 

524 from_masked_image = MaskedImage._read_legacy_hdus( 

525 hdu_list, 

526 filename, 

527 preserve_quantization=preserve_quantization, 

528 plane_map=plane_map, 

529 component=component, 

530 ) 

531 obs_info = _obs_info_from_md(primary_header) 

532 if component is not None: 

533 # This is the image, mask, or variance; attach the projection 

534 # and return 

535 return from_masked_image.view(projection=projection) 

536 result = VisitImage( 

537 from_masked_image.image, 

538 mask=from_masked_image.mask, 

539 variance=from_masked_image.variance, 

540 projection=projection, 

541 psf=psf, 

542 obs_info=obs_info, 

543 ) 

544 result._opaque_metadata = from_masked_image._opaque_metadata 

545 return result 

546 

547 

548class VisitImageSerializationModel[P: pydantic.BaseModel](MaskedImageSerializationModel[P]): 

549 """A Pydantic model used to represent a serialized `VisitImage`.""" 

550 

551 # Inherited attributes are duplicated because that improves the docs 

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

553 # important docs. 

554 

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

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

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

558 ) 

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

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

561 ) 

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

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

564 ) 

565 psf: PiffSerializationModel | PSFExSerializationModel | Any = pydantic.Field( 

566 union_mode="left_to_right", description="PSF model for the image." 

567 ) 

568 obs_info: ObservationInfo = pydantic.Field( 

569 description="Standardized description of visit metadata", 

570 ) 

571 

572 def deserialize_psf(self, archive: InputArchive[Any]) -> PointSpreadFunction | ArchiveReadError: 

573 """Finish deserializing the PSF model, or *return* any exception 

574 raised in the attempt. 

575 """ 

576 try: 

577 match self.psf: 

578 case PiffSerializationModel(): 

579 return PiffWrapper.deserialize(self.psf, archive) 

580 case PSFExSerializationModel(): 

581 return PSFExWrapper.deserialize(self.psf, archive) 

582 case _: 

583 raise ArchiveReadError("PSF model type not recognized.") 

584 except ArchiveReadError as err: 

585 return err 

586 

587 

588def _extract_or_check_header[T]( 

589 key: str, given_value: T | None, header: Any, coerce: Callable[[Any], T] 

590) -> T: 

591 hdr_value: T | None = None 

592 if (hdr_raw_value := header.get(key)) is not None: 

593 hdr_value = coerce(hdr_raw_value) 

594 if given_value is not None: 

595 if hdr_value is not None and hdr_value != given_value: 

596 raise ValueError( 

597 f"Given value {given_value!r} does not match {hdr_value!r} from header key {key}." 

598 ) 

599 return given_value 

600 else: 

601 if hdr_value is None: 

602 raise ValueError(f"No FITS header key {key} found.") 

603 return hdr_value