Coverage for python/lsst/images/

1# This file is part of lsst-images.

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.

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__ = ("Image", "ImageSerializationModel")

16from collections.abc import Callable, Sequence

17from contextlib import ExitStack

18from types import EllipsisType

19from typing import TYPE_CHECKING, Any, ClassVar, final

21import astropy.io.fits

22import astropy.units

23import astropy.wcs

24import numpy as np

25import numpy.typing as npt

26import pydantic

28from lsst.resources import ResourcePath, ResourcePathExpression

30from . import fits

31from ._generalized_image import GeneralizedImage

32from ._geom import YX, Box

33from ._transforms import Frame, SkyProjection, SkyProjectionSerializationModel

34from .serialization import (

35 ArchiveTree,

36 ArrayReferenceModel,

37 ArrayReferenceQuantityModel,

38 InlineArrayModel,

39 InlineArrayQuantityModel,

40 InputArchive,

41 InvalidParameterError,

42 MetadataValue,

43 OutputArchive,

44 no_header_updates,

45)

46from .utils import is_none

48if TYPE_CHECKING:

49 try:

50 from lsst.afw.image import Image as LegacyImage

51 except ImportError:

52 type LegacyImage = Any # type: ignore[no-redef]

55@final

56class Image(GeneralizedImage):

57 """A 2-d array that may be augmented with units and a nonzero origin.

59 Parameters

60 ----------

61 array_or_fill

62 Array or fill value for the image. If a fill value, ``bbox`` or

63 ``shape`` must be provided.

64 bbox

65 Bounding box for the image.

66 yx0

67 Logical coordinates of the first pixel in the array, ordered ``y``,

68 ``x`` (unless an `XY` instance is passed). Ignored if

69 ``bbox`` is provided. Defaults to zeros.

70 shape

71 Leading dimensions of the array, ordered ``y``, ``x`` (unless an `XY`

72 instance is passed). Only needed if ``array_or_fill`` is not an

73 array and ``bbox`` is not provided. Like the bbox, this does not

74 include the last dimension of the array.

75 dtype

76 Pixel data type override.

77 unit

78 Units for the image's pixel values.

79 sky_projection

80 Projection that maps the pixel grid to the sky.

81 metadata

82 Arbitrary flexible metadata to associate with the image.

84 Notes

85 -----

86 Indexing the `array` attribute of an `Image` does not take into account its

87 ``yx0`` offset, but accessing a subimage by indexing an `Image` with a

88 `Box` does, and the `bbox` of the subimage is set to match its location

89 within the original image.

91 Indexed assignment to a subimage requires consistency between the

92 coordinate systems and units of both operands, but it will automatically

93 select a subimage of the right-hand side and convert compatible units when

94 possible. In other words::

96 a[box] = b

98 is a shortcut for

100 a[box].quantity = b[box].quantity

101

102 An ellipsis (``...``) can be used instead of a `Box` to assign to the full

103 image.

104 """

105

106 def __init__(

107 self,

108 array_or_fill: np.ndarray | int | float = 0,

109 /,

110 *,

111 bbox: Box | None = None,

112 yx0: Sequence[int] | None = None,

113 shape: Sequence[int] | None = None,

114 dtype: npt.DTypeLike | None = None,

115 unit: astropy.units.UnitBase | None = None,

116 sky_projection: SkyProjection[Any] | None = None,

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

118 ):

119 super().__init__(metadata)

120 if isinstance(array_or_fill, np.ndarray):

121 if dtype is not None:

122 array = np.array(array_or_fill, dtype=dtype, copy=None)

123 else:

124 array = array_or_fill

125 if bbox is None:

126 bbox = Box.from_shape(array.shape, start=yx0)

127 elif bbox.shape != array.shape:

128 raise ValueError(

129 f"Explicit bbox shape {bbox.shape} does not match array with shape {array.shape}."

130 )

131 if shape is not None and shape != array.shape:

132 raise ValueError(f"Explicit shape {shape} does not match array with shape {array.shape}.")

133 else:

134 if bbox is None:

135 if shape is None:

136 raise TypeError("No bbox, shape, or array provided.")

137 bbox = Box.from_shape(shape, start=yx0)

138 elif shape is not None and shape != bbox.shape:

139 raise ValueError(f"Explicit shape {shape} does not match bbox shape {bbox.shape}.")

140 array = np.full(bbox.shape, array_or_fill, dtype=dtype)

141 self._array: np.ndarray = array

142 self._bbox: Box = bbox

143 self._unit = unit

144 self._sky_projection = sky_projection

145

146 @property

147 def array(self) -> np.ndarray:

148 """The low-level array (`numpy.ndarray`).

149

150 Assigning to this attribute modifies the existing array in place; the

151 bounding box and underlying data pointer are never changed.

152 """

153 return self._array

154

155 @array.setter

156 def array(self, value: np.ndarray | int | float) -> None:

157 self._array[...] = value

158

159 @property

160 def quantity(self) -> astropy.units.Quantity:

161 """The low-level array with units (`astropy.units.Quantity`).

162

163 Assigning to this attribute modifies the existing array in place; the

164 bounding box and underlying data pointer are never changed.

165 """

166 return astropy.units.Quantity(self._array, self._unit, copy=False)

167

168 @quantity.setter

169 def quantity(self, value: astropy.units.Quantity) -> None:

170 self.quantity[...] = value

171

172 @property

173 def bbox(self) -> Box:

174 """Bounding box for the image (`Box`)."""

175 return self._bbox

176

177 @property

178 def unit(self) -> astropy.units.UnitBase | None:

179 """Units for the image's pixel values (`astropy.units.Unit` or

180 `None`).

181 """

182 return self._unit

183

184 @property

185 def sky_projection(self) -> SkyProjection[Any] | None:

186 """The projection that maps this image's pixel grid to the sky

187 (`SkyProjection` | `None`).

188

189 Notes

190 -----

191 The pixel coordinates used by this projection account for the bounding

192 box ``start``; they are not just array indices.

193 """

194 return self._sky_projection

195

196 def __getitem__(self, bbox: Box | EllipsisType) -> Image:

197 if bbox is ...:

198 return self

199 super().__getitem__(bbox)

200 indices = bbox.slice_within(self._bbox)

201 return self._transfer_metadata(

202 Image(self._array[indices], bbox=bbox, unit=self._unit, sky_projection=self._sky_projection),

203 bbox=bbox,

204 )

205

206 def __setitem__(self, bbox: Box | EllipsisType, value: Image) -> None:

207 self[bbox].quantity[...] = value.quantity

208

209 def __str__(self) -> str:

210 return f"Image({self.bbox!s}, {self.array.dtype.type.__name__})"

211

212 def __repr__(self) -> str:

213 return f"Image(..., bbox={self.bbox!r}, dtype={self.array.dtype!r})"

214

215 def __eq__(self, other: object) -> bool:

216 if not isinstance(other, Image):

217 return NotImplemented

218 return (

219 self._bbox == other._bbox

220 and self._unit == other._unit

221 and np.array_equal(self._array, other._array, equal_nan=True)

222 )

223

224 def copy(self) -> Image:

225 return self._transfer_metadata(

226 Image(self._array.copy(), bbox=self._bbox, unit=self._unit, sky_projection=self._sky_projection),

227 copy=True,

228 )

229

230 def view(

231 self,

232 *,

233 unit: astropy.units.UnitBase | None | EllipsisType = ...,

234 sky_projection: SkyProjection | None | EllipsisType = ...,

235 yx0: Sequence[int] | EllipsisType = ...,

236 ) -> Image:

237 """Make a view of the image, with optional updates."""

238 if unit is ...:

239 unit = self._unit

240 if sky_projection is ...:

241 sky_projection = self._sky_projection

242 if yx0 is ...:

243 yx0 = self._bbox.start

244 return self._transfer_metadata(Image(self._array, yx0=yx0, unit=unit, sky_projection=sky_projection))

245

246 def serialize[P: pydantic.BaseModel](

247 self,

248 archive: OutputArchive[P],

249 *,

250 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,

251 save_projection: bool = True,

252 add_offset_wcs: str | None = "A",

253 tile_shape: tuple[int, ...] | None = None,

254 options_name: str | None = None,

255 ) -> ImageSerializationModel[P]:

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

257

258 Parameters

259 ----------

260 archive

261 Archive to write to.

262 update_header

263 A callback that will be given the FITS header for the HDU

264 containing this image in order to add keys to it. This callback

265 may be provided but will not be called if the output format is not

266 FITS.

267 save_projection

268 If `True`, save the `SkyProjection` attached to the image, if there

269 is one. This does not affect whether a FITS WCS corresponding to

270 the projection is written (it always is, if available, and if

271 ``add_offset_wcs`` is not ``" "``).

272 add_offset_wcs

273 A FITS WCS single-character suffix to use when adding a linear

274 WCS that maps the FITS array to the logical pixel coordinates

275 defined by ``bbox.start``. Set to `None` to not write this WCS.

276 If this is set to ``" "``, it will prevent the `SkyProjection` from

277 being saved as a FITS WCS.

278 tile_shape

279 The recommended shape of each tile, if the archive will save

280 the array in distinct tiles for faster subarray retrieval.

281 This is a hint; archives are not required to use this value.

282 options_name

283 Use this name to look up archive options.

284 """

285

286 def _update_header(header: astropy.io.fits.Header) -> None:

287 update_header(header)

288 if self.unit is not None:

289 try:

290 header["BUNIT"] = self.unit.to_string(format="fits")

291 except ValueError:

292 # Units not supported by FITS; write it anyway because

293 # the accepted units are just a recommendation in the

294 # standard.

295 header["BUNIT"] = self.unit.to_string()

296 if self.sky_projection is not None and add_offset_wcs != " ":

297 if self.fits_wcs:

298 header.update(self.fits_wcs.to_header(relax=True))

299 if add_offset_wcs is not None:

300 fits.add_offset_wcs(header, x=self.bbox.x.start, y=self.bbox.y.start, key=add_offset_wcs)

301

302 array_model = archive.add_array(

303 self.array, update_header=_update_header, tile_shape=tile_shape, options_name=options_name

304 )

305 serialized_projection: SkyProjectionSerializationModel[P] | None = None

306 if save_projection and self.sky_projection is not None:

307 serialized_projection = archive.serialize_direct("sky_projection", self.sky_projection.serialize)

308 data = array_model if self.unit is None else array_model.with_units(self.unit)

309 return ImageSerializationModel.model_construct(

310 data=data,

311 yx0=list(self.bbox.start),

312 sky_projection=serialized_projection,

313 metadata=self.metadata,

314 )

315

316 @staticmethod

317 def _get_archive_tree_type[P: pydantic.BaseModel](

318 pointer_type: type[P],

319 ) -> type[ImageSerializationModel[P]]:

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

321 type that uses the given pointer type.

322 """

323 return ImageSerializationModel[pointer_type] # type: ignore

324

325 _archive_default_name: ClassVar[str] = "image"

326 """The name this object should be serialized with when written as the

327 top-level object.

328 """

329

330 @staticmethod

331 def from_legacy(legacy: LegacyImage, unit: astropy.units.UnitBase | None = None) -> Image:

332 """Convert from an `lsst.afw.image.Image` instance.

333

334 Parameters

335 ----------

336 legacy

337 An `lsst.afw.image.Image` instance that will share pixel data with

338 the returned object.

339 unit

340 Units of the image.

341 """

342 return Image(legacy.array, yx0=YX(y=legacy.getY0(), x=legacy.getX0()), unit=unit)

343

344 def to_legacy(self, *, copy: bool | None = None) -> LegacyImage:

345 """Convert to an `lsst.afw.image.Image` instance.

346

347 Parameters

348 ----------

349 copy

350 If `True`, always copy the pixel data. If `False`, return a view,

351 and raise `TypeError` if the pixel data is read-only (this is not

352 supported by afw). If `None`, onyl if the pixel data is

353 read-only.

354 """

355 import lsst.afw.image

356 import lsst.geom

357

358 array = self._array

359 if copy:

360 array = array.copy()

361 elif not self._array.flags.writeable:

362 if copy is None:

363 array = array.copy()

364 else:

365 raise TypeError("Cannot create a legacy lsst.afw.image.Image view into a read-only array.")

366

367 return lsst.afw.image.Image(

368 array,

369 deep=False,

370 dtype=array.dtype.type,

371 xy0=lsst.geom.Point2I(self._bbox.x.min, self._bbox.y.min),

372 )

373

374 @staticmethod

375 def read_legacy(

376 uri: ResourcePathExpression,

377 *,

378 preserve_quantization: bool = False,

379 ext: str | int = 1,

380 fits_wcs_frame: Frame | None = None,

381 ) -> Image:

382 """Read a FITS file written by `lsst.afw.image.Image.writeFits`.

383

384 Parameters

385 ----------

386 uri

387 URI or file name.

388 preserve_quantization

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

390 exactly preserve quantization-compressed pixel values. This causes

391 the arrays to be marked as read-only and stores the original binary

392 table data for those planes in memory. If the `Image` is copied,

393 the precompressed pixel values are not transferred to the copy.

394 ext

395 Name or index of the FITS HDU to read.

396 fits_wcs_frame

397 If not `None` and the HDU containing the image has a FITS WCS,

398 attach a `SkyProjection` to the returned image by converting that

399 WCS.

400 """

401 opaque_metadata = fits.FitsOpaqueMetadata()

402 with ExitStack() as exit_stack:

403 fs, fspath = ResourcePath(uri).to_fsspec()

404 stream = exit_stack.enter_context(fs.open(fspath))

405 hdu_list = exit_stack.enter_context(astropy.io.fits.open(stream))

406 opaque_metadata.extract_legacy_primary_header(hdu_list[0].header)

407 bintable_hdu: astropy.io.fits.BinTableHDU | None = None

408 if preserve_quantization:

409 bintable_stream = exit_stack.enter_context(fs.open(fspath))

410 bintable_hdu_list = exit_stack.enter_context(

411 astropy.io.fits.open(bintable_stream, disable_image_compression=True)

412 )

413 bintable_hdu = bintable_hdu_list[ext]

414 result = Image._read_legacy_hdu(

415 hdu_list[ext], opaque_metadata, preserve_bintable=bintable_hdu, fits_wcs_frame=fits_wcs_frame

416 )

417 result._opaque_metadata = opaque_metadata

418 return result

419

420 @staticmethod

421 def _read_legacy_hdu(

422 hdu: astropy.io.fits.ImageHDU | astropy.io.fits.CompImageHDU,

423 opaque_metadata: fits.FitsOpaqueMetadata,

424 *,

425 preserve_bintable: astropy.io.fits.BinTableHDU | None,

426 fits_wcs_frame: Frame | None = None,

427 ) -> Image:

428 unit: astropy.units.UnitBase | None = None

429 if (fits_unit := hdu.header.pop("BUNIT", None)) is not None:

430 try:

431 unit = astropy.units.Unit(fits_unit, format="fits")

432 except ValueError:

433 # Accept non-FITS units by assuming Astropy can still figure

434 # them out if we don't specify the format.

435 unit = astropy.units.Unit(fits_unit)

436 if opaque_metadata.get_instrumental_unit() == astropy.units.electron:

437 # Fix incorrect BUNIT='adu' in LSST preliminary_visit_image.

438 if unit == astropy.units.adu:

439 unit = astropy.units.electron

440 if unit == astropy.units.adu**2:

441 unit = astropy.units.electron**2

442 dx: int = hdu.header.pop("LTV1")

443 dy: int = hdu.header.pop("LTV2")

444 yx0 = YX(y=-dy, x=-dx)

445 read_only: bool = False

446 if preserve_bintable is not None:

447 opaque_metadata.precompressed[hdu.name] = fits.PrecompressedImage.from_bintable(preserve_bintable)

448 read_only = True

449 sky_projection: SkyProjection | None = None

450 if fits_wcs_frame is not None:

451 try:

452 fits_wcs = astropy.wcs.WCS(hdu.header)

453 except KeyError:

454 pass

455 else:

456 sky_projection = SkyProjection.from_fits_wcs(

457 fits_wcs, pixel_frame=fits_wcs_frame, x0=yx0.x, y0=yx0.y

458 )

459 image = Image(hdu.data, yx0=yx0, unit=unit, sky_projection=sky_projection)

460 if read_only:

461 image._array.flags["WRITEABLE"] = False

462 fits.strip_wcs_cards(hdu.header)

463 hdu.header.strip()

464 hdu.header.remove("EXTTYPE", ignore_missing=True)

465 hdu.header.remove("INHERIT", ignore_missing=True)

466 hdu.header.remove("UZSCALE", ignore_missing=True)

467 opaque_metadata.add_header(hdu.header)

468 return image

469

470

471class ImageSerializationModel[P: pydantic.BaseModel](ArchiveTree):

472 """Pydantic model used to represent the serialized form of an `.Image`."""

473

474 SCHEMA_NAME: ClassVar[str] = "image"

475 SCHEMA_VERSION: ClassVar[str] = "1.0.0"

476 MIN_READ_VERSION: ClassVar[int] = 1

477 PUBLIC_TYPE: ClassVar[type] = Image

478

479 data: ArrayReferenceQuantityModel | ArrayReferenceModel | InlineArrayModel | InlineArrayQuantityModel = (

480 pydantic.Field(description="Reference to pixel data.")

481 )

482 yx0: list[int] = pydantic.Field(

483 description="Coordinate of the first pixels in the array, ordered (y, x)."

484 )

485 sky_projection: SkyProjectionSerializationModel[P] | None = pydantic.Field(

486 default=None,

487 exclude_if=is_none,

488 description="Projection that maps the logical pixel grid onto the sky.",

489 )

490

491 @property

492 def bbox(self) -> Box:

493 """The bounding box of the image."""

494 match self.data:

495 case ArrayReferenceQuantityModel() | InlineArrayQuantityModel():

496 shape = self.data.value.shape

497 case ArrayReferenceModel() | InlineArrayModel():

498 shape = self.data.shape

499 return Box.from_shape(shape, self.yx0)

500

501 def deserialize(

502 self,

503 archive: InputArchive[Any],

504 *,

505 bbox: Box | None = None,

506 strip_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,

507 **kwargs: Any,

508 ) -> Image:

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

510

511 Parameters

512 ----------

513 archive

514 Archive to read from.

515 bbox

516 Bounding box of a subimage to read instead.

517 strip_header

518 A callable that strips out any FITS header cards added by the

519 ``update_header`` argument in the corresponding call to

520 `Image.serialize`.

521 **kwargs

522 Unsupported keyword arguments are accepted only to provide better

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

524 """

525 if kwargs:

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

527 array_model: ArrayReferenceModel | InlineArrayModel

528 unit: astropy.units.UnitBase | None = None

529 if isinstance(self.data, ArrayReferenceQuantityModel | InlineArrayQuantityModel):

530 array_model = self.data.value

531 unit = self.data.unit

532 else:

533 array_model = self.data

534

535 def _strip_header(header: astropy.io.fits.Header) -> None:

536 if unit is not None:

537 header.pop("BUNIT", None)

538 fits.strip_wcs_cards(header)

539 strip_header(header)

540

541 slices = bbox.slice_within(self.bbox) if bbox is not None else ...

542 array = archive.get_array(array_model, strip_header=_strip_header, slices=slices)

543 sky_projection = self.sky_projection.deserialize(archive) if self.sky_projection is not None else None

544 return Image(

545 array,

546 yx0=self.yx0 if bbox is None else bbox.start,

547 unit=unit,

548 sky_projection=sky_projection,

549 )._finish_deserialize(self)

550

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

552 if kwargs:

553 raise InvalidParameterError(f"Unsupported parameters for Image components: {set(kwargs.keys())}.")

554 return super().deserialize_component(component, archive)

Coverage for python/lsst/images/_image.py: 23%

219 statements