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

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 

14import warnings 

15 

16__all__ = ("VisitImage", "VisitImageSerializationModel") 

17 

18import functools 

19from collections.abc import Callable, Mapping, MutableMapping 

20from types import EllipsisType 

21from typing import Any, Literal, cast, overload 

22 

23import astropy.io.fits 

24import astropy.units 

25import astropy.wcs 

26import pydantic 

27from astro_metadata_translator import ObservationInfo, VisitInfoTranslator 

28 

29from ._geom import Box 

30from ._image import Image, ImageSerializationModel 

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

32from ._masked_image import MaskedImage, MaskedImageSerializationModel 

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

34from .fits import FitsOpaqueMetadata 

35from .psfs import ( 

36 PiffSerializationModel, 

37 PiffWrapper, 

38 PointSpreadFunction, 

39 PSFExSerializationModel, 

40 PSFExWrapper, 

41) 

42from .serialization import ArchiveReadError, InputArchive, OutputArchive 

43 

44 

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

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

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

48 try: 

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

50 except ValueError: 

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

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

53 # complete. 

54 if visit_info is not None: 

55 from lsst.afw.image import setVisitInfoMetadata 

56 from lsst.daf.base import PropertyList 

57 

58 pl = PropertyList() 

59 setVisitInfoMetadata(pl, visit_info) 

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

61 md.update(pl) 

62 

63 # Try the given header looking for VisitInfo hints. 

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

65 # no way to disable those without capturing them. 

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

67 return obs_info 

68 

69 

70class VisitImage(MaskedImage): 

71 """A calibrated single-visit image. 

72 

73 Parameters 

74 ---------- 

75 image 

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

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

78 mask 

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

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

81 is replaced (possibly by `None`). 

82 variance 

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

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

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

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

87 (possibly by `None`). 

88 mask_schema 

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

90 not provided. 

91 opaque_metadata 

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

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

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

95 projection 

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

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

98 psf 

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

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

101 obs_info 

102 General information about this visit in standardized form. 

103 """ 

104 

105 def __init__( 

106 self, 

107 image: Image, 

108 *, 

109 mask: Mask | None = None, 

110 variance: Image | None = None, 

111 mask_schema: MaskSchema | None = None, 

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

113 obs_info: ObservationInfo | None = None, 

114 psf: PointSpreadFunction | ArchiveReadError, 

115 ): 

116 super().__init__( 

117 image, 

118 mask=mask, 

119 variance=variance, 

120 mask_schema=mask_schema, 

121 projection=projection, 

122 obs_info=obs_info, 

123 ) 

124 if self.image.unit is None: 

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

126 if self.image.projection is None: 

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

128 if self.image.obs_info is None: 

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

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

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

132 self._psf = psf 

133 

134 @property 

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

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

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

138 

139 @property 

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

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

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

143 """ 

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

145 

146 @property 

147 def obs_info(self) -> ObservationInfo: 

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

149 (`~astro_metadata_translator.ObservationInfo`). 

150 """ 

151 obs_info = self.image.obs_info 

152 assert obs_info is not None 

153 return obs_info 

154 

155 @property 

156 def astropy_wcs(self) -> ProjectionAstropyView: 

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

158 

159 Notes 

160 ----- 

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

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

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

164 

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

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

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

168 """ 

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

170 

171 @property 

172 def psf(self) -> PointSpreadFunction: 

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

174 (`.psfs.PointSpreadFunction`). 

175 """ 

176 if isinstance(self._psf, ArchiveReadError): 

177 raise self._psf 

178 return self._psf 

179 

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

181 super().__getitem__(bbox) 

182 if bbox is ...: 

183 return self 

184 result = VisitImage( 

185 self.image[bbox], 

186 mask=self.mask[bbox], 

187 variance=self.variance[bbox], 

188 projection=self.projection, 

189 psf=self.psf, 

190 obs_info=self.obs_info, 

191 ) 

192 if opaque_metadata := getattr(self, "_opaque_metadata", None): 

193 result._opaque_metadata = opaque_metadata.subset(bbox) 

194 return result 

195 

196 def __str__(self) -> str: 

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

198 

199 def __repr__(self) -> str: 

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

201 

202 def copy(self) -> VisitImage: 

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

204 result = VisitImage( 

205 image=self._image.copy(), 

206 mask=self._mask.copy(), 

207 variance=self._variance.copy(), 

208 psf=self._psf, 

209 obs_info=self.obs_info, 

210 ) 

211 if self._opaque_metadata is not None: 

212 result._opaque_metadata = self._opaque_metadata.copy() 

213 return result 

214 

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

216 serialized_image = archive.serialize_direct( 

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

218 ) 

219 serialized_mask = archive.serialize_direct( 

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

221 ) 

222 serialized_variance = archive.serialize_direct( 

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

224 ) 

225 serialized_projection = ( 

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

227 if self.projection is not None 

228 else None 

229 ) 

230 serialized_psf: PiffSerializationModel | PSFExSerializationModel 

231 match self._psf: 

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

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

234 case PiffWrapper(): 

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

236 case PSFExWrapper(): 

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

238 case _: 

239 raise TypeError( 

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

241 ) 

242 return VisitImageSerializationModel( 

243 image=serialized_image, 

244 mask=serialized_mask, 

245 variance=serialized_variance, 

246 projection=serialized_projection, 

247 psf=serialized_psf, 

248 obs_info=self.obs_info, 

249 ) 

250 

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

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

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

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

255 # and here we do not want it to. 

256 @staticmethod 

257 def deserialize( 

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

259 archive: InputArchive[Any], 

260 *, 

261 bbox: Box | None = None, 

262 ) -> VisitImage: 

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

264 psf = model.deserialize_psf(archive) 

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

266 result = VisitImage( 

267 masked_image.image, 

268 mask=masked_image.mask, 

269 variance=masked_image.variance, 

270 psf=psf, 

271 projection=projection, 

272 obs_info=model.obs_info, 

273 ) 

274 return result 

275 

276 @staticmethod 

277 def _get_archive_tree_type[P: pydantic.BaseModel]( 

278 pointer_type: type[P], 

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

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

281 type that uses the given pointer type. 

282 """ 

283 return VisitImageSerializationModel[pointer_type] # type: ignore 

284 

285 # write_fits and read_fits inherited from MaskedImage. 

286 

287 @staticmethod 

288 def from_legacy( 

289 legacy: Any, 

290 *, 

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

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

293 instrument: str | None = None, 

294 visit: int | None = None, 

295 component: str | None = None, 

296 ) -> VisitImage: 

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

298 

299 Parameters 

300 ---------- 

301 legacy 

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

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

304 unit 

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

306 key will be used, if available. 

307 plane_map 

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

309 description. If `None` (default) 

310 `get_legacy_visit_image_mask_planes` is used. 

311 instrument 

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

313 provided. 

314 visit 

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

316 """ 

317 if plane_map is None: 

318 plane_map = get_legacy_visit_image_mask_planes() 

319 md = legacy.getMetadata() 

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

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

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

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

324 legacy_wcs = legacy.getWcs() 

325 if legacy_wcs is None: 

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

327 legacy_detector = legacy.getDetector() 

328 if legacy_detector is None: 

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

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

331 opaque_fits_metadata = FitsOpaqueMetadata() 

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

333 with warnings.catch_warnings(): 

334 # Silence warnings about long keys becoming HIERARCH. 

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

336 primary_header.update(md.toOrderedDict()) 

337 opaque_fits_metadata.extract_legacy_primary_header(primary_header) 

338 projection = Projection.from_legacy( 

339 legacy_wcs, 

340 DetectorFrame( 

341 instrument=instrument, 

342 visit=visit, 

343 detector=legacy_detector.getId(), 

344 bbox=detector_bbox, 

345 ), 

346 ) 

347 legacy_psf = legacy.getPsf() 

348 if legacy_psf is None: 

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

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

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

352 result = VisitImage( 

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

354 mask=masked_image.mask, 

355 variance=masked_image.variance, 

356 projection=projection, 

357 psf=psf, 

358 obs_info=obs_info, 

359 ) 

360 result._opaque_metadata = masked_image._opaque_metadata 

361 return result 

362 

363 @overload # type: ignore[override] 

364 @staticmethod 

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

366 filename: str, 

367 *, 

368 component: Literal["bbox"], 

369 ) -> Box: ... 

370 

371 @overload 

372 @staticmethod 

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

374 filename: str, 

375 *, 

376 preserve_quantization: bool = False, 

377 instrument: str | None = None, 

378 visit: int | None = None, 

379 component: Literal["image"], 

380 ) -> Image: ... 

381 

382 @overload 

383 @staticmethod 

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

385 filename: str, 

386 *, 

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

388 instrument: str | None = None, 

389 visit: int | None = None, 

390 component: Literal["mask"], 

391 ) -> Mask: ... 

392 

393 @overload 

394 @staticmethod 

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

396 filename: str, 

397 *, 

398 preserve_quantization: bool = False, 

399 instrument: str | None = None, 

400 visit: int | None = None, 

401 component: Literal["variance"], 

402 ) -> Image: ... 

403 

404 @overload 

405 @staticmethod 

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

407 filename: str, 

408 *, 

409 instrument: str | None = None, 

410 visit: int | None = None, 

411 component: Literal["projection"], 

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

413 

414 @overload 

415 @staticmethod 

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

417 filename: str, 

418 *, 

419 component: Literal["psf"], 

420 ) -> PointSpreadFunction: ... 

421 

422 @overload 

423 @staticmethod 

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

425 filename: str, 

426 *, 

427 preserve_quantization: bool = False, 

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

429 instrument: str | None = None, 

430 visit: int | None = None, 

431 component: None = None, 

432 ) -> VisitImage: ... 

433 

434 @staticmethod 

435 def read_legacy( # type: ignore[override] 

436 filename: str, 

437 *, 

438 preserve_quantization: bool = False, 

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

440 instrument: str | None = None, 

441 visit: int | None = None, 

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

443 ) -> Any: 

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

445 

446 Parameters 

447 ---------- 

448 filename 

449 Full name of the file. 

450 preserve_quantization 

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

452 exactly preserve quantization-compressed pixel values. This causes 

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

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

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

456 not transferred to the copy. 

457 plane_map 

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

459 description. If `None` (default) 

460 `get_legacy_visit_image_mask_planes` is used. 

461 instrument 

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

463 provided. 

464 visit 

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

466 provided. 

467 component 

468 A component to read instead of the full image. 

469 """ 

470 from lsst.afw.image import ExposureFitsReader 

471 

472 reader = ExposureFitsReader(filename) 

473 if component == "bbox": 

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

475 legacy_detector = reader.readDetector() 

476 if legacy_detector is None: 

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

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

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

480 legacy_wcs = reader.readWcs() 

481 if legacy_wcs is None: 

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

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

484 legacy_psf = reader.readPsf() 

485 if legacy_psf is None: 

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

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

488 if component == "psf": 

489 return psf 

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

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

492 primary_header = hdu_list[0].header 

493 if instrument is None: 

494 try: 

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

496 except LookupError: 

497 raise ValueError( 

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

499 ) from None 

500 if visit is None: 

501 try: 

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

503 except LookupError: 

504 raise ValueError( 

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

506 ) from None 

507 projection = Projection.from_legacy( 

508 legacy_wcs, 

509 DetectorFrame( 

510 instrument=instrument, 

511 visit=visit, 

512 detector=legacy_detector.getId(), 

513 bbox=detector_bbox, 

514 ), 

515 ) 

516 if component == "projection": 

517 return projection 

518 if plane_map is None: 

519 plane_map = get_legacy_visit_image_mask_planes() 

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

521 from_masked_image = MaskedImage._read_legacy_hdus( 

522 hdu_list, 

523 filename, 

524 preserve_quantization=preserve_quantization, 

525 plane_map=plane_map, 

526 component=component, 

527 ) 

528 obs_info = _obs_info_from_md(primary_header) 

529 if component is not None: 

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

531 # and return 

532 return from_masked_image.view(projection=projection) 

533 result = VisitImage( 

534 from_masked_image.image, 

535 mask=from_masked_image.mask, 

536 variance=from_masked_image.variance, 

537 projection=projection, 

538 psf=psf, 

539 obs_info=obs_info, 

540 ) 

541 result._opaque_metadata = from_masked_image._opaque_metadata 

542 return result 

543 

544 

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

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

547 

548 # Inherited attributes are duplicated because that improves the docs 

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

550 # important docs. 

551 

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

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

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

555 ) 

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

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

558 ) 

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

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

561 ) 

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

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

564 ) 

565 obs_info: ObservationInfo = pydantic.Field( 

566 description="Standardized description of visit metadata", 

567 ) 

568 

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

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

571 raised in the attempt. 

572 """ 

573 try: 

574 match self.psf: 

575 case PiffSerializationModel(): 

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

577 case PSFExSerializationModel(): 

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

579 case _: 

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

581 except ArchiveReadError as err: 

582 return err 

583 

584 

585def _extract_or_check_header[T]( 

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

587) -> T: 

588 hdr_value: T | None = None 

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

590 hdr_value = coerce(hdr_raw_value) 

591 if given_value is not None: 

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

593 raise ValueError( 

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

595 ) 

596 return given_value 

597 else: 

598 if hdr_value is None: 

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

600 return hdr_value