Coverage for python/lsst/images/fits/_common.py: 87%
223 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-06-24 08:46 +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.
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 return cls(str(row["EXTNAME"]).upper(), int(row["EXTVER"]))
105 @classmethod
106 def from_str(cls, source: str) -> ExtensionKey:
107 """Construct from the `str` coercion of this type, which is used
108 as the 'source' field in various Pydantic models that serve as
109 references to other HDUs.
110 """
111 if (m := FITS_SOURCE_REGEX.fullmatch(source)) is None: 111 ↛ 112line 111 didn't jump to line 112 because the condition on line 111 was never true
112 raise ArchiveReadError(f"Bad 'source' string for FITS: {source!r}.")
113 ver = 1
114 if m.group("extver") is not None:
115 ver = int(m.group("extver"))
116 return cls(m.group("extname"), ver)
118 def check(self) -> None:
119 if not FITS_SOURCE_REGEX.match(str(self)): 119 ↛ 120line 119 didn't jump to line 120 because the condition on line 119 was never true
120 raise ValueError(
121 f"Invalid source key: '{str(self)}'; name characters must be alphanumeric, '-', '_', or '/'."
122 )
124 def __str__(self) -> str:
125 if self.ver > 1:
126 return f"fits:{self.name},{self.ver}"
127 else:
128 return f"fits:{self.name}"
131class PointerModel(pydantic.BaseModel):
132 column: TableColumnModel = pydantic.Field(description="Table column for this cell this pointer targets.")
133 row: int = pydantic.Field(description="Zero-indexed row for cell this pointer targets.")
136class InvalidFitsArchiveError(RuntimeError):
137 """The error type raised when the content of a FITS file presumed to have
138 been written by FitsOutputArchive is not self-consistent.
139 """
142class FitsCompressionAlgorithm(enum.StrEnum):
143 """FITS compression algorithms supported by this package.
145 See the FITS standard for definitions.
146 """
148 GZIP_1 = "GZIP_1"
149 GZIP_2 = "GZIP_2"
150 RICE_1 = "RICE_1"
153class FitsDitherAlgorithm(enum.StrEnum):
154 """FITS quantization dither algorithms supported by this package.
156 See the FITS standard for definitions.
157 """
159 NO_DITHER = "NO_DITHER"
160 SUBTRACTIVE_DITHER_1 = "SUBTRACTIVE_DITHER_1"
161 SUBTRACTIVE_DITHER_2 = "SUBTRACTIVE_DITHER_2"
163 def to_astropy_quantize_method(self) -> int:
164 """Convert to the integer code used by Astropy."""
165 match self:
166 case self.NO_DITHER: 166 ↛ 167line 166 didn't jump to line 167 because the pattern on line 166 never matched
167 return -1
168 case self.SUBTRACTIVE_DITHER_1: 168 ↛ 169line 168 didn't jump to line 169 because the pattern on line 168 never matched
169 return 1
170 case self.SUBTRACTIVE_DITHER_2:
171 return 2
172 raise AssertionError("Invalid enum value.")
175class FitsQuantizationOptions(pydantic.BaseModel, frozen=True):
176 """Quantization options for FITS compression."""
178 dither: FitsDitherAlgorithm
179 """How to add random noise during quantization to reduce biases."""
181 level: float
182 """Quantization level.
184 When positive, this is the fraction of the measured standard deviation that
185 corresponds to an integer step. When negative, it is ``-ZSCALE``, the
186 scaling to apply directly to the original pixels before quantization.
187 """
189 seed: int | None = None
190 """Random number seed to use for dithering.
192 Values between 1 and 10000 (inclusive) are used directly. ``0`` will
193 generate a value from the current time, and ``-1`` will generate a value
194 from the checksum of the image.
196 If `None`, the ``compression_seed`` parameter must be passed to
197 `FitsOutputArchive.open` if any quantized compression is configured.
198 """
201class FitsCompressionOptions(pydantic.BaseModel, frozen=True):
202 """Configuration options for FITS compression."""
204 algorithm: FitsCompressionAlgorithm = FitsCompressionAlgorithm.GZIP_2
205 """Compression algorithm to use."""
207 tile_shape: tuple[int, ...] | None = None
208 """Shape ``(..., y, x)`` of independently compressed tiles.
210 The default of `None` leaves the tile shape up to the ``tile_shape``
211 argument to `~lsst.images.serialization.OutputArchive.add_array`.
212 """
214 quantization: FitsQuantizationOptions | None = None
215 """Quantization to apply before compression, if any."""
217 DEFAULT: ClassVar[FitsCompressionOptions | None]
218 """Default compression options (lossless ``GZIP_2``)."""
220 LOSSY: ClassVar[FitsCompressionOptions]
221 """Default lossy compression options."""
223 def make_hdu(self, data: np.ndarray, name: str) -> astropy.io.fits.CompImageHDU:
224 """Make an `astropy.io.fits.CompImageHDU` object from these options."""
225 if self.quantization is not None:
226 return astropy.io.fits.CompImageHDU(
227 data,
228 name=name,
229 compression_type=self.algorithm.value,
230 tile_shape=self.tile_shape,
231 quantize_method=self.quantization.dither.to_astropy_quantize_method(),
232 quantize_level=self.quantization.level,
233 dither_seed=self.quantization.seed,
234 )
235 else:
236 return astropy.io.fits.CompImageHDU(
237 data,
238 name=name,
239 compression_type=self.algorithm.value,
240 tile_shape=self.tile_shape,
241 quantize_level=0.0,
242 )
245FitsCompressionOptions.DEFAULT = FitsCompressionOptions()
246FitsCompressionOptions.LOSSY = FitsCompressionOptions(
247 algorithm=FitsCompressionAlgorithm.RICE_1,
248 quantization=FitsQuantizationOptions(dither=FitsDitherAlgorithm.SUBTRACTIVE_DITHER_2, level=16.0),
249)
252_COMPRESSION_KEYS = frozenset(
253 (
254 "ZIMAGE",
255 "ZCMPTYPE",
256 "ZBITPIX",
257 "ZNAXIS",
258 "ZMASKCMP",
259 "ZQUANTIZ",
260 "ZDITHER0",
261 "ZCHECKSUM",
262 "ZDATASUM",
263 )
264)
265_COMPRESSION_PREFIX_KEYS = ("ZNAXIS", "ZTILE", "ZNAME", "ZVAL")
268@dataclasses.dataclass
269class PrecompressedImage:
270 """Already-compressed FITS HDUs that are attached to high-level objects
271 via `FitsOpaqueMetadata`, allowing lossy-compressed pixel values to be
272 round-tripped exactly.
273 """
275 header: astropy.io.fits.Header
276 """Header for the HDU.
278 This contains only FITS tile-compression keywords.
279 """
281 data: astropy.io.fits.FITS_rec
282 """FITS binary table data that serves as the low-level representation of a
283 tile-compressed image HDU.
284 """
286 @classmethod
287 def from_bintable(cls, hdu: astropy.io.fits.BinTableHDU) -> Self:
288 """Construct from a binary table HDU.
290 Parameters
291 ----------
292 hdu
293 Binary table HDU, typically read from a FITS file opened with
294 ``disable_image_compression=True``.
296 Returns
297 -------
298 PrecompressedImage
299 A `PrecompressedImage` instance.
300 """
301 header = astropy.io.fits.Header(
302 [
303 card
304 for card in hdu.header.cards
305 if card.keyword in _COMPRESSION_KEYS
306 or any(card.keyword.startswith(k) for k in _COMPRESSION_PREFIX_KEYS)
307 ]
308 )
309 # This is an opportunity to fix CFITSIO's non-standard writing of the
310 # old RICE_ONE value instead of RICE_1.
311 if header["ZCMPTYPE"] == "RICE_ONE":
312 header["ZCMPTYPE"] = "RICE_1"
313 return cls(header=header, data=hdu.data)
316@final
317@dataclasses.dataclass
318class FitsOpaqueMetadata(OpaqueArchiveMetadata):
319 """Opaque metadata that may be carried around by a serializable type to
320 propagate serialization options and opaque information without that type
321 knowing how it was serialized.
322 """
324 headers: dict[ExtensionKey, astropy.io.fits.Header] = dataclasses.field(default_factory=dict)
325 """FITS headers found (but not interpreted and stripped) when reading, to
326 be propagated on write.
328 Keys are EXTNAME/EXTVER combinations, or ("", 1) for the primary header.
329 Header information in opaque metadata is considered immutable, allowing it
330 to be transferred by reference to copies and subsets of the object it is
331 attached to.
332 """
334 precompressed: dict[str, PrecompressedImage] = dataclasses.field(default_factory=dict)
335 """FITS tile-compressed HDUs that should be written out directly instead
336 of the in-memory data provided.
338 Keys are EXTNAME values, which must be unique (no EXTVER disambiguation).
339 Precompressed pixel values are never copied or transferred to subsets.
340 """
342 def add_header(
343 self,
344 header: astropy.io.fits.Header,
345 name: str | None = None,
346 ver: int | None = None,
347 key: ExtensionKey | None = None,
348 ) -> None:
349 """Add a header to the opaque metadata if it is not already present,
350 and strip EXTNAME, EXTVER, and DATE if present.
352 Parameters
353 ----------
354 header
355 Header to add. May be modified in place.
356 name
357 EXTNAME (all caps). If not provided, the EXTNAME card must be
358 present in the header. Use ``""`` for the primary header.
359 ver
360 EXTVER. If not provided and the EXTVER card is not present,
361 defaults to 1.
362 key
363 Combination of EXTNAME and EXTVER; used instead of both ``name``
364 and ``ver`` if provided.
365 """
366 if key is None:
367 if name is None:
368 name = header["EXTNAME"]
369 if ver is None:
370 ver = header.get("EXTVER", 1)
371 key = ExtensionKey(name, ver)
372 strip_butler_cards(header)
373 if key not in self.headers: 373 ↛ exitline 373 didn't return from function 'add_header' because the condition on line 373 was always true
374 header.remove("EXTNAME", ignore_missing=True)
375 header.remove("EXTVER", ignore_missing=True)
376 # DATE is written fresh on every HDU, so it is regenerated on write
377 # rather than propagated as opaque metadata.
378 header.remove("DATE", ignore_missing=True)
379 self.headers[key] = header
381 def maybe_use_precompressed(self, name: str) -> astropy.io.fits.BinTableHDU | None:
382 """Look up the given EXTNAME to see if there is a tile compressed image
383 HDU that should be used directly, instead of requantizing.
385 Parameters
386 ----------
387 name
388 EXTNAME (all caps).
390 Returns
391 -------
392 `astropy.io.fits.BinTableHDU` | `None`
393 An already-compressed HDU, in binary table form, or `None` if there
394 is no precompressed HDU for this EXTNAME.
395 """
396 if (precompressed := self.precompressed.get(name)) is None: 396 ↛ 398line 396 didn't jump to line 398 because the condition on line 396 was always true
397 return None
398 return astropy.io.fits.BinTableHDU(precompressed.data, header=precompressed.header.copy(), name=name)
400 def extract_legacy_primary_header(self, header: astropy.io.fits.Header) -> dict[str, Any]:
401 """Update the opaque metadata with the header of the primary HDU
402 of a legacy (`lsst.afw.image`) FITS file, stripping cards we know we
403 don't need and extracting any ``LSST IMAGES ...`` cards into a
404 dictionary we return.
405 """
406 primary_header = header.copy(strip=True)
407 # No idea what these spare TAN-SIP headers are doing in the afw
408 # FITS files, but we'll strip them here:
409 primary_header.remove("A_ORDER", ignore_missing=True)
410 primary_header.remove("B_ORDER", ignore_missing=True)
411 primary_header.remove("DATE", ignore_missing=True)
412 strip_legacy_exposure_cards(primary_header)
413 strip_butler_cards(primary_header)
414 metadata: dict[str, Any] = {}
415 for n in itertools.count(): 415 ↛ 420line 415 didn't jump to line 420 because the loop on line 415 didn't complete
416 if (key := header.pop(f"LSST IMAGES KEY {n + 1}", ...)) is ...: 416 ↛ 418line 416 didn't jump to line 418 because the condition on line 416 was always true
417 break
418 value = header.pop(f"LSST IMAGES VALUE {n + 1}")
419 metadata[key] = value
420 self.headers[ExtensionKey()] = primary_header
421 return metadata
423 def add_cutdown_primary_header(self, header: astropy.io.fits.Header) -> None:
424 """Add the primary header of a cut-down ``lsst.images`` FITS file as
425 opaque metadata.
427 A cut-down file is one whose JSON-tree, index, and nested-archive HDUs
428 have been dropped, leaving only the primary, image, mask, and variance
429 HDUs (e.g. as written by ``dax_images_cutout``). Only the
430 container-layout cards that would be misleading after re-serialization
431 are stripped; all other cards are retained so they survive the round
432 trip.
434 Parameters
435 ----------
436 header
437 Primary header to read. Not modified.
438 """
439 primary_header = header.copy(strip=True)
440 for keyword in ("FMTVER", "INDXADDR", "INDXSIZE", "JSONADDR", "JSONSIZE", "DATAMODL"):
441 primary_header.remove(keyword, ignore_missing=True)
442 self.add_header(primary_header, name="", ver=1)
444 def copy(self) -> FitsOpaqueMetadata:
445 # Docstring inherited.
446 return FitsOpaqueMetadata(headers=self.headers)
448 def subset(self, bbox: Box) -> FitsOpaqueMetadata:
449 # Docstring inherited.
450 return FitsOpaqueMetadata(headers=self.headers)
452 def get_instrumental_unit(self) -> astropy.units.UnitBase | None:
453 """Extract the ``LSST ISR UNIT`` key from the primary header (if it
454 exists) and wrap it with Astropy.
455 """
456 if (primary_header := self.headers.get(ExtensionKey())) is not None: 456 ↛ 459line 456 didn't jump to line 459 because the condition on line 456 was always true
457 if (instrumental_unit_str := primary_header.get("LSST ISR UNITS")) is not None: 457 ↛ 458line 457 didn't jump to line 458 because the condition on line 457 was never true
458 return astropy.units.Unit(instrumental_unit_str)
459 return None
462def add_offset_wcs(header: astropy.io.fits.Header, *, x: int | float, y: int | float, key: str = "A") -> None:
463 """Add a trivial FITS WCS to a header that applies the appropriate offset
464 to map FITS array coordinates to a logical pixel grid.
466 Parameters
467 ----------
468 header
469 Header to update in-place.
470 x
471 Logical coordinate of the first column.
472 y
473 Logical coordinate of the first row.
474 key
475 Single-character suffix for this WCS.
476 """
477 header.set(f"CTYPE1{key}", "LINEAR")
478 header.set(f"CTYPE2{key}", "LINEAR")
479 header.set(f"CRPIX1{key}", 1.0)
480 header.set(f"CRPIX2{key}", 1.0)
481 header.set(f"CRVAL1{key}", float(x))
482 header.set(f"CRVAL2{key}", float(y))
483 header.set(f"CUNIT1{key}", "PIXEL")
484 header.set(f"CUNIT2{key}", "PIXEL")
487def read_offset_wcs(header: astropy.io.fits.Header, *, key: str = "A") -> tuple[int, int] | None:
488 """Recover the logical pixel origin from a linear offset WCS written by
489 `add_offset_wcs`.
491 Parameters
492 ----------
493 header
494 Header to read from.
495 key
496 Single-character suffix for the WCS to read.
498 Returns
499 -------
500 `tuple` [`int`, `int`] | `None`
501 The ``(x, y)`` logical coordinates of the first column and row, or
502 `None` if the header has no ``LINEAR`` WCS with this suffix.
503 """
504 if header.get(f"CTYPE1{key}") != "LINEAR":
505 return None
506 x = header[f"CRVAL1{key}"] - (header.get(f"CRPIX1{key}", 1.0) - 1.0)
507 y = header[f"CRVAL2{key}"] - (header.get(f"CRPIX2{key}", 1.0) - 1.0)
508 return (round(x), round(y))
511def read_yx0(header: astropy.io.fits.Header) -> YX[int]:
512 """Recover the logical origin of an array from its FITS header.
514 Parameters
515 ----------
516 header
517 Header to read from. Not modified.
519 Returns
520 -------
521 `~lsst.images.YX` [`int`]
522 Logical coordinate of the first pixel, ordered ``(y, x)``.
524 Raises
525 ------
526 ValueError
527 Raised if the header records no origin via either the ``LTV1``/``LTV2``
528 cards (written by `lsst.afw.image`) or a linear offset WCS (written by
529 `add_offset_wcs`).
531 Notes
532 -----
533 The ``LTV1``/``LTV2`` convention is preferred when present, falling back to
534 the ``A`` offset WCS otherwise.
535 """
536 if "LTV1" in header:
537 return YX(y=-round(header["LTV2"]), x=-round(header["LTV1"]))
538 if (xy0 := read_offset_wcs(header)) is not None:
539 return YX(y=xy0[1], x=xy0[0])
540 raise ValueError("Header records no LTV1/LTV2 cards or linear offset WCS.")
543_WCS_VECTOR_KEYS = ("CUNIT", "CRPIX", "CRPIX", "CRVAL", "CRDELT", "CROTA", "CRDER", "CSYER", "CDELT")
544_WCS_MATRIX_KEYS = ("CD{0}_{1}", "PC{0}_{1}")
547def strip_wcs_cards(header: astropy.io.fits.Header) -> None:
548 """Strip WCS cards from a FITS header.
550 This does *not* attempt to cover all possible FITS WCS forms; it focuses on
551 the ones we actually plan to write (simple undistorted ones + TAN-SIP).
552 """
553 wcsaxes = header.pop("WCSAXES", 2)
554 for wcsname in [""] + list(string.ascii_uppercase):
555 header.remove("RADESYS" + wcsname, ignore_missing=True)
556 if "CTYPE1" + wcsname in header:
557 ctype: str = "" # just for linters that can't figure out that the loop always executes
558 for n in range(wcsaxes):
559 suffix = f"{n + 1}{wcsname}"
560 ctype = header.pop("CTYPE" + suffix)
561 for key in _WCS_VECTOR_KEYS:
562 header.remove(key + suffix, ignore_missing=True)
563 for m in range(wcsaxes):
564 for tmpl in _WCS_MATRIX_KEYS:
565 header.remove(tmpl.format(m + 1, suffix), ignore_missing=True)
566 if ctype.endswith("-SIP"): 566 ↛ 567line 566 didn't jump to line 567 because the condition on line 566 was never true
567 _strip_sip_poly(header, wcsname, "A")
568 _strip_sip_poly(header, wcsname, "B")
569 _strip_sip_poly(header, wcsname, "AP")
570 _strip_sip_poly(header, wcsname, "BP")
571 header.remove("LONPOLE", ignore_missing=True)
572 header.remove("LATPOLE", ignore_missing=True)
573 header.remove("MJDREF", ignore_missing=True)
574 header.remove("PV1_3", ignore_missing=True)
575 header.remove("PV1_4", ignore_missing=True)
578def _strip_sip_poly(header: astropy.io.fits.Header, wcsname: str, which: str) -> None:
579 order: int | None = header.pop(f"{which}_ORDER{wcsname}", None)
580 if order is not None:
581 for i, j in itertools.product(range(order + 1), range(order + 1)):
582 header.remove(f"{which}_{i}_{j}{wcsname}", ignore_missing=True)
585def strip_legacy_exposure_cards(header: astropy.io.fits.Header) -> None:
586 """Strip header keywords added by lsst.afw.image.Exposure."""
587 header.remove("AR_HDU", ignore_missing=True)
588 for name in (
589 "FILTER",
590 "DETECTOR",
591 "VALID_POLYGON",
592 "SKYWCS",
593 "PSF",
594 "SUMMARYSTATS",
595 "AP_CORR_MAP",
596 "PHOTOCALIB",
597 ):
598 header.remove(f"{name}_ID", ignore_missing=True)
599 header.remove(f"ARCHIVE_ID_{name}", ignore_missing=True)
602def strip_butler_cards(header: astropy.io.fits.Header) -> None:
603 """Strip header keywords added by butler provenance that would be
604 incorrect if propagated to a downstream file.
605 """
606 for key in list(header):
607 if key.startswith("LSST BUTLER"):
608 del header[key]