Coverage for python/lsst/images/fits/

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.

13from __future__ import annotations

15__all__ = (

16 "FITS_SOURCE_REGEX",

17 "JSON_COLUMN",

18 "JSON_EXTNAME",

19 "ExtensionHDU",

20 "ExtensionKey",

21 "FitsCompressionAlgorithm",

22 "FitsCompressionOptions",

23 "FitsDitherAlgorithm",

24 "FitsOpaqueMetadata",

25 "FitsQuantizationOptions",

26 "InvalidFitsArchiveError",

27 "PointerModel",

28 "PrecompressedImage",

29 "add_offset_wcs",

30 "read_offset_wcs",

31 "read_yx0",

32 "strip_butler_cards",

33 "strip_legacy_exposure_cards",

34 "strip_wcs_cards",

35 "suppress_fits_card_warnings",

36)

38import contextlib

39import dataclasses

40import enum

41import itertools

42import re

43import string

44import warnings

45from collections.abc import Iterator

46from typing import Any, ClassVar, Self, final

48import astropy.io.fits

49import astropy.io.fits.verify

50import numpy as np

51import pydantic

53from .._geom import YX, Box

54from ..serialization import ArchiveReadError, OpaqueArchiveMetadata, TableColumnModel

56type ExtensionHDU = astropy.io.fits.ImageHDU | astropy.io.fits.CompImageHDU | astropy.io.fits.BinTableHDU

58FITS_SOURCE_REGEX = re.compile(r"fits:(?P<extname>[\w/\-]+)(,(?P<extver>\d+))?(\[\d+\])?")

60JSON_EXTNAME: str = "JSON"

61JSON_COLUMN: str = "JSON"

64@contextlib.contextmanager

65def suppress_fits_card_warnings() -> Iterator[None]:

66 """Silence Astropy warnings for FITS cards it fixes up automatically.

68 When a header is written, Astropy converts keywords longer than eight

69 characters into ``HIERARCH`` cards and truncates comments that overflow a

70 card. Both emit a `~astropy.io.fits.verify.VerifyWarning` that carries no

71 actionable information, so they are silenced around the code that writes

72 these headers.

73 """

74 with warnings.catch_warnings():

75 warnings.filterwarnings(

76 "ignore",

77 message="Keyword name .* is greater than 8 characters",

78 category=astropy.io.fits.verify.VerifyWarning,

79 )

80 warnings.filterwarnings(

81 "ignore",

82 message=".*comment will be truncated",

83 category=astropy.io.fits.verify.VerifyWarning,

84 )

85 yield

88@dataclasses.dataclass(frozen=True)

89class ExtensionKey:

90 """The identifiers for a single FITS extension HDU within a file."""

92 name: str = ""

93 """Value of the EXTNAME keyword, or an empty string for the primary HDU."""

95 ver: int = 1

96 """Value of the EXTVER keyword (which may be absent if it is ``1``)."""

98 @classmethod

99 def from_index_row(cls, row: np.ndarray) -> ExtensionKey:

100 """Construct from a row of the index binary table appended to the

101 FITS files written by the package.

102

103 Parameters

104 ----------

105 row

106 Structured-array row with ``EXTNAME`` and ``EXTVER`` fields.

107 """

108 return cls(str(row["EXTNAME"]).upper(), int(row["EXTVER"]))

109

110 @classmethod

111 def from_str(cls, source: str) -> ExtensionKey:

112 """Construct from the `str` coercion of this type, which is used

113 as the 'source' field in various Pydantic models that serve as

114 references to other HDUs.

115

116 Parameters

117 ----------

118 source

119 ``fits:<extname>[,<extver>]`` reference string, as produced by

120 ``str(self)``.

121 """

122 if (m := FITS_SOURCE_REGEX.fullmatch(source)) is None: 122 ↛ 123line 122 didn't jump to line 123 because the condition on line 122 was never true

123 raise ArchiveReadError(f"Bad 'source' string for FITS: {source!r}.")

124 ver = 1

125 if m.group("extver") is not None:

126 ver = int(m.group("extver"))

127 return cls(m.group("extname"), ver)

128

129 def check(self) -> None:

130 if not FITS_SOURCE_REGEX.match(str(self)): 130 ↛ 131line 130 didn't jump to line 131 because the condition on line 130 was never true

131 raise ValueError(

132 f"Invalid source key: '{str(self)}'; name characters must be alphanumeric, '-', '_', or '/'."

133 )

134

135 def __str__(self) -> str:

136 if self.ver > 1:

137 return f"fits:{self.name},{self.ver}"

138 else:

139 return f"fits:{self.name}"

140

141

142class PointerModel(pydantic.BaseModel):

143 column: TableColumnModel = pydantic.Field(description="Table column for this cell this pointer targets.")

144 row: int = pydantic.Field(description="Zero-indexed row for cell this pointer targets.")

145

146

147class InvalidFitsArchiveError(RuntimeError):

148 """The error type raised when the content of a FITS file presumed to have

149 been written by FitsOutputArchive is not self-consistent.

150 """

151

152

153class FitsCompressionAlgorithm(enum.StrEnum):

154 """FITS compression algorithms supported by this package.

155

156 See the FITS standard for definitions.

157 """

158

159 GZIP_1 = "GZIP_1"

160 GZIP_2 = "GZIP_2"

161 RICE_1 = "RICE_1"

162

163

164class FitsDitherAlgorithm(enum.StrEnum):

165 """FITS quantization dither algorithms supported by this package.

166

167 See the FITS standard for definitions.

168 """

169

170 NO_DITHER = "NO_DITHER"

171 SUBTRACTIVE_DITHER_1 = "SUBTRACTIVE_DITHER_1"

172 SUBTRACTIVE_DITHER_2 = "SUBTRACTIVE_DITHER_2"

173

174 def to_astropy_quantize_method(self) -> int:

175 """Convert to the integer code used by Astropy."""

176 match self:

177 case self.NO_DITHER: 177 ↛ 178line 177 didn't jump to line 178 because the pattern on line 177 never matched

178 return -1

179 case self.SUBTRACTIVE_DITHER_1: 179 ↛ 180line 179 didn't jump to line 180 because the pattern on line 179 never matched

180 return 1

181 case self.SUBTRACTIVE_DITHER_2:

182 return 2

183 raise AssertionError("Invalid enum value.")

184

185

186class FitsQuantizationOptions(pydantic.BaseModel, frozen=True):

187 """Quantization options for FITS compression."""

188

189 dither: FitsDitherAlgorithm

190 """How to add random noise during quantization to reduce biases."""

191

192 level: float

193 """Quantization level.

194

195 When positive, this is the fraction of the measured standard deviation that

196 corresponds to an integer step. When negative, it is ``-ZSCALE``, the

197 scaling to apply directly to the original pixels before quantization.

198 """

199

200 seed: int | None = None

201 """Random number seed to use for dithering.

202

203 Values between 1 and 10000 (inclusive) are used directly. ``0`` will

204 generate a value from the current time, and ``-1`` will generate a value

205 from the checksum of the image.

206

207 If `None`, the ``compression_seed`` parameter must be passed to

208 `FitsOutputArchive.open` if any quantized compression is configured.

209 """

210

211

212class FitsCompressionOptions(pydantic.BaseModel, frozen=True):

213 """Configuration options for FITS compression."""

214

215 algorithm: FitsCompressionAlgorithm = FitsCompressionAlgorithm.GZIP_2

216 """Compression algorithm to use."""

217

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

219 """Shape ``(..., y, x)`` of independently compressed tiles.

220

221 The default of `None` leaves the tile shape up to the ``tile_shape``

222 argument to `~lsst.images.serialization.OutputArchive.add_array`.

223 """

224

225 quantization: FitsQuantizationOptions | None = None

226 """Quantization to apply before compression, if any."""

227

228 DEFAULT: ClassVar[FitsCompressionOptions | None]

229 """Default compression options (lossless ``GZIP_2``)."""

230

231 LOSSY: ClassVar[FitsCompressionOptions]

232 """Default lossy compression options."""

233

234 def make_hdu(self, data: np.ndarray, name: str) -> astropy.io.fits.CompImageHDU:

235 """Make an `astropy.io.fits.CompImageHDU` object from these options.

236

237 Parameters

238 ----------

239 data

240 Pixel data to store in the HDU.

241 name

242 ``EXTNAME`` of the HDU.

243 """

244 if self.quantization is not None:

245 return astropy.io.fits.CompImageHDU(

246 data,

247 name=name,

248 compression_type=self.algorithm.value,

249 tile_shape=self.tile_shape,

250 quantize_method=self.quantization.dither.to_astropy_quantize_method(),

251 quantize_level=self.quantization.level,

252 dither_seed=self.quantization.seed,

253 )

254 else:

255 return astropy.io.fits.CompImageHDU(

256 data,

257 name=name,

258 compression_type=self.algorithm.value,

259 tile_shape=self.tile_shape,

260 quantize_level=0.0,

261 )

262

263

264FitsCompressionOptions.DEFAULT = FitsCompressionOptions()

265FitsCompressionOptions.LOSSY = FitsCompressionOptions(

266 algorithm=FitsCompressionAlgorithm.RICE_1,

267 quantization=FitsQuantizationOptions(dither=FitsDitherAlgorithm.SUBTRACTIVE_DITHER_2, level=16.0),

268)

269

270

271_COMPRESSION_KEYS = frozenset(

272 (

273 "ZIMAGE",

274 "ZCMPTYPE",

275 "ZBITPIX",

276 "ZNAXIS",

277 "ZMASKCMP",

278 "ZQUANTIZ",

279 "ZDITHER0",

280 "ZCHECKSUM",

281 "ZDATASUM",

282 )

283)

284_COMPRESSION_PREFIX_KEYS = ("ZNAXIS", "ZTILE", "ZNAME", "ZVAL")

285

286

287@dataclasses.dataclass

288class PrecompressedImage:

289 """Already-compressed FITS HDUs that are attached to high-level objects

290 via `FitsOpaqueMetadata`, allowing lossy-compressed pixel values to be

291 round-tripped exactly.

292 """

293

294 header: astropy.io.fits.Header

295 """Header for the HDU.

296

297 This contains only FITS tile-compression keywords.

298 """

299

300 data: astropy.io.fits.FITS_rec

301 """FITS binary table data that serves as the low-level representation of a

302 tile-compressed image HDU.

303 """

304

305 @classmethod

306 def from_bintable(cls, hdu: astropy.io.fits.BinTableHDU) -> Self:

307 """Construct from a binary table HDU.

308

309 Parameters

310 ----------

311 hdu

312 Binary table HDU, typically read from a FITS file opened with

313 ``disable_image_compression=True``.

314

315 Returns

316 -------

317 PrecompressedImage

318 A `PrecompressedImage` instance.

319 """

320 header = astropy.io.fits.Header(

321 [

322 card

323 for card in hdu.header.cards

324 if card.keyword in _COMPRESSION_KEYS

325 or any(card.keyword.startswith(k) for k in _COMPRESSION_PREFIX_KEYS)

326 ]

327 )

328 # This is an opportunity to fix CFITSIO's non-standard writing of the

329 # old RICE_ONE value instead of RICE_1.

330 if header["ZCMPTYPE"] == "RICE_ONE":

331 header["ZCMPTYPE"] = "RICE_1"

332 return cls(header=header, data=hdu.data)

333

334

335@final

336@dataclasses.dataclass

337class FitsOpaqueMetadata(OpaqueArchiveMetadata):

338 """Opaque metadata that may be carried around by a serializable type to

339 propagate serialization options and opaque information without that type

340 knowing how it was serialized.

341 """

342

343 headers: dict[ExtensionKey, astropy.io.fits.Header] = dataclasses.field(default_factory=dict)

344 """FITS headers found (but not interpreted and stripped) when reading, to

345 be propagated on write.

346

347 Keys are EXTNAME/EXTVER combinations, or ("", 1) for the primary header.

348 Header information in opaque metadata is considered immutable, allowing it

349 to be transferred by reference to copies and subsets of the object it is

350 attached to.

351 """

352

353 precompressed: dict[str, PrecompressedImage] = dataclasses.field(default_factory=dict)

354 """FITS tile-compressed HDUs that should be written out directly instead

355 of the in-memory data provided.

356

357 Keys are EXTNAME values, which must be unique (no EXTVER disambiguation).

358 Precompressed pixel values are never copied or transferred to subsets.

359 """

360

361 def add_header(

362 self,

363 header: astropy.io.fits.Header,

364 name: str | None = None,

365 ver: int | None = None,

366 key: ExtensionKey | None = None,

367 ) -> None:

368 """Add a header to the opaque metadata if it is not already present,

369 and strip EXTNAME, EXTVER, and DATE if present.

370

371 Parameters

372 ----------

373 header

374 Header to add. May be modified in place.

375 name

376 EXTNAME (all caps). If not provided, the EXTNAME card must be

377 present in the header. Use ``""`` for the primary header.

378 ver

379 EXTVER. If not provided and the EXTVER card is not present,

380 defaults to 1.

381 key

382 Combination of EXTNAME and EXTVER; used instead of both ``name``

383 and ``ver`` if provided.

384 """

385 if key is None:

386 if name is None:

387 name = header["EXTNAME"]

388 if ver is None:

389 ver = header.get("EXTVER", 1)

390 key = ExtensionKey(name, ver)

391 strip_butler_cards(header)

392 if key not in self.headers: 392 ↛ exitline 392 didn't return from function 'add_header' because the condition on line 392 was always true

393 header.remove("EXTNAME", ignore_missing=True)

394 header.remove("EXTVER", ignore_missing=True)

395 # DATE is written fresh on every HDU, so it is regenerated on write

396 # rather than propagated as opaque metadata.

397 header.remove("DATE", ignore_missing=True)

398 self.headers[key] = header

399

400 def maybe_use_precompressed(self, name: str) -> astropy.io.fits.BinTableHDU | None:

401 """Look up the given EXTNAME to see if there is a tile compressed image

402 HDU that should be used directly, instead of requantizing.

403

404 Parameters

405 ----------

406 name

407 EXTNAME (all caps).

408

409 Returns

410 -------

411 `astropy.io.fits.BinTableHDU` | `None`

412 An already-compressed HDU, in binary table form, or `None` if there

413 is no precompressed HDU for this EXTNAME.

414 """

415 if (precompressed := self.precompressed.get(name)) is None: 415 ↛ 417line 415 didn't jump to line 417 because the condition on line 415 was always true

416 return None

417 return astropy.io.fits.BinTableHDU(precompressed.data, header=precompressed.header.copy(), name=name)

418

419 def extract_legacy_primary_header(self, header: astropy.io.fits.Header) -> dict[str, Any]:

420 """Update the opaque metadata with the header of the primary HDU

421 of a legacy (`lsst.afw.image`) FITS file, stripping cards we know we

422 don't need and extracting any ``LSST IMAGES ...`` cards into a

423 dictionary we return.

424

425 Parameters

426 ----------

427 header

428 Primary HDU header of the legacy FITS file.

429 """

430 primary_header = header.copy(strip=True)

431 # No idea what these spare TAN-SIP headers are doing in the afw

432 # FITS files, but we'll strip them here:

433 primary_header.remove("A_ORDER", ignore_missing=True)

434 primary_header.remove("B_ORDER", ignore_missing=True)

435 primary_header.remove("DATE", ignore_missing=True)

436 strip_legacy_exposure_cards(primary_header)

437 strip_butler_cards(primary_header)

438 metadata: dict[str, Any] = {}

439 for n in itertools.count(): 439 ↛ 444line 439 didn't jump to line 444 because the loop on line 439 didn't complete

440 if (key := header.pop(f"LSST IMAGES KEY {n + 1}", ...)) is ...: 440 ↛ 442line 440 didn't jump to line 442 because the condition on line 440 was always true

441 break

442 value = header.pop(f"LSST IMAGES VALUE {n + 1}")

443 metadata[key] = value

444 self.headers[ExtensionKey()] = primary_header

445 return metadata

446

447 def add_cutdown_primary_header(self, header: astropy.io.fits.Header) -> None:

448 """Add the primary header of a cut-down ``lsst.images`` FITS file as

449 opaque metadata.

450

451 A cut-down file is one whose JSON-tree, index, and nested-archive HDUs

452 have been dropped, leaving only the primary, image, mask, and variance

453 HDUs (e.g. as written by ``dax_images_cutout``). Only the

454 container-layout cards that would be misleading after re-serialization

455 are stripped; all other cards are retained so they survive the round

456 trip.

457

458 Parameters

459 ----------

460 header

461 Primary header to read. Not modified.

462 """

463 primary_header = header.copy(strip=True)

464 for keyword in ("FMTVER", "INDXADDR", "INDXSIZE", "JSONADDR", "JSONSIZE", "DATAMODL"):

465 primary_header.remove(keyword, ignore_missing=True)

466 self.add_header(primary_header, name="", ver=1)

467

468 def copy(self) -> FitsOpaqueMetadata:

469 # Docstring inherited.

470 return FitsOpaqueMetadata(headers=self.headers)

471

472 def subset(self, bbox: Box) -> FitsOpaqueMetadata:

473 # Docstring inherited.

474 return FitsOpaqueMetadata(headers=self.headers)

475

476 def get_instrumental_unit(self) -> astropy.units.UnitBase | None:

477 """Extract the ``LSST ISR UNIT`` key from the primary header (if it

478 exists) and wrap it with Astropy.

479 """

480 if (primary_header := self.headers.get(ExtensionKey())) is not None: 480 ↛ 483line 480 didn't jump to line 483 because the condition on line 480 was always true

481 if (instrumental_unit_str := primary_header.get("LSST ISR UNITS")) is not None: 481 ↛ 482line 481 didn't jump to line 482 because the condition on line 481 was never true

482 return astropy.units.Unit(instrumental_unit_str)

483 return None

484

485

486def add_offset_wcs(header: astropy.io.fits.Header, *, x: int | float, y: int | float, key: str = "A") -> None:

487 """Add a trivial FITS WCS to a header that applies the appropriate offset

488 to map FITS array coordinates to a logical pixel grid.

489

490 Parameters

491 ----------

492 header

493 Header to update in-place.

494 x

495 Logical coordinate of the first column.

496 y

497 Logical coordinate of the first row.

498 key

499 Single-character suffix for this WCS.

500 """

501 header.set(f"CTYPE1{key}", "LINEAR")

502 header.set(f"CTYPE2{key}", "LINEAR")

503 header.set(f"CRPIX1{key}", 1.0)

504 header.set(f"CRPIX2{key}", 1.0)

505 header.set(f"CRVAL1{key}", float(x))

506 header.set(f"CRVAL2{key}", float(y))

507 header.set(f"CUNIT1{key}", "PIXEL")

508 header.set(f"CUNIT2{key}", "PIXEL")

509

510

511def read_offset_wcs(header: astropy.io.fits.Header, *, key: str = "A") -> tuple[int, int] | None:

512 """Recover the logical pixel origin from a linear offset WCS written by

513 `add_offset_wcs`.

514

515 Parameters

516 ----------

517 header

518 Header to read from.

519 key

520 Single-character suffix for the WCS to read.

521

522 Returns

523 -------

524 `tuple` [`int`, `int`] | `None`

525 The ``(x, y)`` logical coordinates of the first column and row, or

526 `None` if the header has no ``LINEAR`` WCS with this suffix.

527 """

528 if header.get(f"CTYPE1{key}") != "LINEAR":

529 return None

530 x = header[f"CRVAL1{key}"] - (header.get(f"CRPIX1{key}", 1.0) - 1.0)

531 y = header[f"CRVAL2{key}"] - (header.get(f"CRPIX2{key}", 1.0) - 1.0)

532 return (round(x), round(y))

533

534

535def read_yx0(header: astropy.io.fits.Header) -> YX[int]:

536 """Recover the logical origin of an array from its FITS header.

537

538 Parameters

539 ----------

540 header

541 Header to read from. Not modified.

542

543 Returns

544 -------

545 `~lsst.images.YX` [`int`]

546 Logical coordinate of the first pixel, ordered ``(y, x)``.

547

548 Raises

549 ------

550 ValueError

551 Raised if the header records no origin via either the ``LTV1``/``LTV2``

552 cards (written by `lsst.afw.image`) or a linear offset WCS (written by

553 `add_offset_wcs`).

554

555 Notes

556 -----

557 The ``LTV1``/``LTV2`` convention is preferred when present, falling back to

558 the ``A`` offset WCS otherwise.

559 """

560 if "LTV1" in header:

561 return YX(y=-round(header["LTV2"]), x=-round(header["LTV1"]))

562 if (xy0 := read_offset_wcs(header)) is not None:

563 return YX(y=xy0[1], x=xy0[0])

564 raise ValueError("Header records no LTV1/LTV2 cards or linear offset WCS.")

565

566

567_WCS_VECTOR_KEYS = ("CUNIT", "CRPIX", "CRPIX", "CRVAL", "CRDELT", "CROTA", "CRDER", "CSYER", "CDELT")

568_WCS_MATRIX_KEYS = ("CD{0}_{1}", "PC{0}_{1}")

569

570

571def strip_wcs_cards(header: astropy.io.fits.Header) -> None:

572 """Strip WCS cards from a FITS header.

573

574 This does *not* attempt to cover all possible FITS WCS forms; it focuses on

575 the ones we actually plan to write (simple undistorted ones + TAN-SIP).

576

577 Parameters

578 ----------

579 header

580 FITS header to strip WCS cards from in place.

581 """

582 wcsaxes = header.pop("WCSAXES", 2)

583 for wcsname in [""] + list(string.ascii_uppercase):

584 header.remove("RADESYS" + wcsname, ignore_missing=True)

585 if "CTYPE1" + wcsname in header:

586 ctype: str = "" # just for linters that can't figure out that the loop always executes

587 for n in range(wcsaxes):

588 suffix = f"{n + 1}{wcsname}"

589 ctype = header.pop("CTYPE" + suffix)

590 for key in _WCS_VECTOR_KEYS:

591 header.remove(key + suffix, ignore_missing=True)

592 for m in range(wcsaxes):

593 for tmpl in _WCS_MATRIX_KEYS:

594 header.remove(tmpl.format(m + 1, suffix), ignore_missing=True)

595 if ctype.endswith("-SIP"): 595 ↛ 596line 595 didn't jump to line 596 because the condition on line 595 was never true

596 _strip_sip_poly(header, wcsname, "A")

597 _strip_sip_poly(header, wcsname, "B")

598 _strip_sip_poly(header, wcsname, "AP")

599 _strip_sip_poly(header, wcsname, "BP")

600 header.remove("LONPOLE", ignore_missing=True)

601 header.remove("LATPOLE", ignore_missing=True)

602 header.remove("MJDREF", ignore_missing=True)

603 header.remove("PV1_3", ignore_missing=True)

604 header.remove("PV1_4", ignore_missing=True)

605

606

607def _strip_sip_poly(header: astropy.io.fits.Header, wcsname: str, which: str) -> None:

608 order: int | None = header.pop(f"{which}_ORDER{wcsname}", None)

609 if order is not None:

610 for i, j in itertools.product(range(order + 1), range(order + 1)):

611 header.remove(f"{which}_{i}_{j}{wcsname}", ignore_missing=True)

612

613

614def strip_legacy_exposure_cards(header: astropy.io.fits.Header) -> None:

615 """Strip header keywords added by lsst.afw.image.Exposure.

616

617 Parameters

618 ----------

619 header

620 FITS header to strip exposure cards from in place.

621 """

622 header.remove("AR_HDU", ignore_missing=True)

623 for name in (

624 "FILTER",

625 "DETECTOR",

626 "VALID_POLYGON",

627 "SKYWCS",

628 "PSF",

629 "SUMMARYSTATS",

630 "AP_CORR_MAP",

631 "PHOTOCALIB",

632 ):

633 header.remove(f"{name}_ID", ignore_missing=True)

634 header.remove(f"ARCHIVE_ID_{name}", ignore_missing=True)

635

636

637def strip_butler_cards(header: astropy.io.fits.Header) -> None:

638 """Strip header keywords added by butler provenance that would be

639 incorrect if propagated to a downstream file.

640

641 Parameters

642 ----------

643 header

644 FITS header to strip butler provenance cards from in place.

645 """

646 for key in list(header):

647 if key.startswith("LSST BUTLER"):

648 del header[key]

Coverage for python/lsst/images/fits/_common.py: 87%

223 statements