Coverage for python/lsst/images/fits/_common.py: 37%
194 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-10 02:26 -0700
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 "strip_butler_cards",
31 "strip_legacy_exposure_cards",
32 "strip_wcs_cards",
33)
35import dataclasses
36import enum
37import itertools
38import re
39import string
40from typing import Any, ClassVar, Self, final
42import astropy.io.fits
43import numpy as np
44import pydantic
46from .._geom import Box
47from ..serialization import ArchiveReadError, OpaqueArchiveMetadata, TableColumnModel
49type ExtensionHDU = astropy.io.fits.ImageHDU | astropy.io.fits.CompImageHDU | astropy.io.fits.BinTableHDU
51FITS_SOURCE_REGEX = re.compile(r"fits:(?P<extname>[\w/\-]+)(,(?P<extver>\d+))?(\[\d+\])?")
53JSON_EXTNAME: str = "JSON"
54JSON_COLUMN: str = "JSON"
57@dataclasses.dataclass(frozen=True)
58class ExtensionKey:
59 """The identifiers for a single FITS extension HDU within a file."""
61 name: str = ""
62 """Value of the EXTNAME keyword, or an empty string for the primary HDU."""
64 ver: int = 1
65 """Value of the EXTVER keyword (which may be absent if it is ``1``)."""
67 @classmethod
68 def from_index_row(cls, row: np.ndarray) -> ExtensionKey:
69 """Construct from a row of the index binary table appended to the
70 FITS files written by the package.
71 """
72 return cls(str(row["EXTNAME"]).upper(), int(row["EXTVER"]))
74 @classmethod
75 def from_str(cls, source: str) -> ExtensionKey:
76 """Construct from the `str` coercion of this type, which is used
77 as the 'source' field in various Pydantic models that serve as
78 references to other HDUs.
79 """
80 if (m := FITS_SOURCE_REGEX.fullmatch(source)) is None:
81 raise ArchiveReadError(f"Bad 'source' string for FITS: {source!r}.")
82 ver = 1
83 if m.group("extver") is not None:
84 ver = int(m.group("extver"))
85 return cls(m.group("extname"), ver)
87 def check(self) -> None:
88 if not FITS_SOURCE_REGEX.match(str(self)):
89 raise ValueError(
90 f"Invalid source key: '{str(self)}'; name characters must be alphanumeric, '-', '_', or '/'."
91 )
93 def __str__(self) -> str:
94 if self.ver > 1:
95 return f"fits:{self.name},{self.ver}"
96 else:
97 return f"fits:{self.name}"
100class PointerModel(pydantic.BaseModel):
101 column: TableColumnModel = pydantic.Field(description="Table column for this cell this pointer targets.")
102 row: int = pydantic.Field(description="Zero-indexed row for cell this pointer targets.")
105class InvalidFitsArchiveError(RuntimeError):
106 """The error type raised when the content of a FITS file presumed to have
107 been written by FitsOutputArchive is not self-consistent.
108 """
111class FitsCompressionAlgorithm(enum.StrEnum):
112 """FITS compression algorithms supported by this package.
114 See the FITS standard for definitions.
115 """
117 GZIP_1 = "GZIP_1"
118 GZIP_2 = "GZIP_2"
119 RICE_1 = "RICE_1"
122class FitsDitherAlgorithm(enum.StrEnum):
123 """FITS quantization dither algorithms supported by this package.
125 See the FITS standard for definitions.
126 """
128 NO_DITHER = "NO_DITHER"
129 SUBTRACTIVE_DITHER_1 = "SUBTRACTIVE_DITHER_1"
130 SUBTRACTIVE_DITHER_2 = "SUBTRACTIVE_DITHER_2"
132 def to_astropy_quantize_method(self) -> int:
133 """Convert to the integer code used by Astropy."""
134 match self:
135 case self.NO_DITHER:
136 return -1
137 case self.SUBTRACTIVE_DITHER_1:
138 return 1
139 case self.SUBTRACTIVE_DITHER_2:
140 return 2
141 raise AssertionError("Invalid enum value.")
144class FitsQuantizationOptions(pydantic.BaseModel, frozen=True):
145 """Quantization options for FITS compression."""
147 dither: FitsDitherAlgorithm
148 """How to add random noise during quantization to reduce biases."""
150 level: float
151 """Quantization level.
153 When positive, this is the fraction of the measured standard deviation that
154 corresponds to an integer step. When negative, it is ``-ZSCALE``, the
155 scaling to apply directly to the original pixels before quantization.
156 """
158 seed: int | None = None
159 """Random number seed to use for dithering.
161 Values between 1 and 10000 (inclusive) are used directly. ``0`` will
162 generate a value from the current time, and ``-1`` will generate a value
163 from the checksum of the image.
165 If `None`, the ``compression_seed`` parameter must be passed to
166 `FitsOutputArchive.open` if any quantized compression is configured.
167 """
170class FitsCompressionOptions(pydantic.BaseModel, frozen=True):
171 """Configuration options for FITS compression."""
173 algorithm: FitsCompressionAlgorithm = FitsCompressionAlgorithm.GZIP_2
174 """Compression algorithm to use."""
176 tile_shape: tuple[int, ...] | None = None
177 """Shape ``(..., y, x)`` of independently compressed tiles.
179 The default of `None` leaves the tile shape up to the ``tile_shape``
180 argument to `~lsst.images.serialization.OutputArchive.add_array`.
181 """
183 quantization: FitsQuantizationOptions | None = None
184 """Quantization to apply before compression, if any."""
186 DEFAULT: ClassVar[FitsCompressionOptions | None]
187 """Default compression options (lossless ``GZIP_2``)."""
189 LOSSY: ClassVar[FitsCompressionOptions]
190 """Default lossy compression options."""
192 def make_hdu(self, data: np.ndarray, name: str) -> astropy.io.fits.CompImageHDU:
193 """Make an `astropy.io.fits.CompImageHDU` object from these options."""
194 if self.quantization is not None:
195 return astropy.io.fits.CompImageHDU(
196 data,
197 name=name,
198 compression_type=self.algorithm.value,
199 tile_shape=self.tile_shape,
200 quantize_method=self.quantization.dither.to_astropy_quantize_method(),
201 quantize_level=self.quantization.level,
202 dither_seed=self.quantization.seed,
203 )
204 else:
205 return astropy.io.fits.CompImageHDU(
206 data,
207 name=name,
208 compression_type=self.algorithm.value,
209 tile_shape=self.tile_shape,
210 quantize_level=0.0,
211 )
214FitsCompressionOptions.DEFAULT = FitsCompressionOptions()
215FitsCompressionOptions.LOSSY = FitsCompressionOptions(
216 algorithm=FitsCompressionAlgorithm.RICE_1,
217 quantization=FitsQuantizationOptions(dither=FitsDitherAlgorithm.SUBTRACTIVE_DITHER_2, level=16.0),
218)
221_COMPRESSION_KEYS = frozenset(
222 (
223 "ZIMAGE",
224 "ZCMPTYPE",
225 "ZBITPIX",
226 "ZNAXIS",
227 "ZMASKCMP",
228 "ZQUANTIZ",
229 "ZDITHER0",
230 "ZCHECKSUM",
231 "ZDATASUM",
232 )
233)
234_COMPRESSION_PREFIX_KEYS = ("ZNAXIS", "ZTILE", "ZNAME", "ZVAL")
237@dataclasses.dataclass
238class PrecompressedImage:
239 """Already-compressed FITS HDUs that are attached to high-level objects
240 via `FitsOpaqueMetadata`, allowing lossy-compressed pixel values to be
241 round-tripped exactly.
242 """
244 header: astropy.io.fits.Header
245 """Header for the HDU.
247 This contains only FITS tile-compression keywords.
248 """
250 data: astropy.io.fits.FITS_rec
251 """FITS binary table data that serves as the low-level representation of a
252 tile-compressed image HDU.
253 """
255 @classmethod
256 def from_bintable(cls, hdu: astropy.io.fits.BinTableHDU) -> Self:
257 """Construct from a binary table HDU.
259 Parameters
260 ----------
261 hdu
262 Binary table HDU, typically read from a FITS file opened with
263 ``disable_image_compression=True``.
265 Returns
266 -------
267 PrecompressedImage
268 A `PrecompressedImage` instance.
269 """
270 header = astropy.io.fits.Header(
271 [
272 card
273 for card in hdu.header.cards
274 if card.keyword in _COMPRESSION_KEYS
275 or any(card.keyword.startswith(k) for k in _COMPRESSION_PREFIX_KEYS)
276 ]
277 )
278 # This is an opportunity to fix CFITSIO's non-standard writing of the
279 # old RICE_ONE value instead of RICE_1.
280 if header["ZCMPTYPE"] == "RICE_ONE":
281 header["ZCMPTYPE"] = "RICE_1"
282 return cls(header=header, data=hdu.data)
285@final
286@dataclasses.dataclass
287class FitsOpaqueMetadata(OpaqueArchiveMetadata):
288 """Opaque metadata that may be carried around by a serializable type to
289 propagate serialization options and opaque information without that type
290 knowing how it was serialized.
291 """
293 headers: dict[ExtensionKey, astropy.io.fits.Header] = dataclasses.field(default_factory=dict)
294 """FITS headers found (but not interpreted and stripped) when reading, to
295 be propagated on write.
297 Keys are EXTNAME/EXTVER combinations, or ("", 1) for the primary header.
298 Header information in opaque metadata is considered immutable, allowing it
299 to be transferred by reference to copies and subsets of the object it is
300 attached to.
301 """
303 precompressed: dict[str, PrecompressedImage] = dataclasses.field(default_factory=dict)
304 """FITS tile-compressed HDUs that should be written out directly instead
305 of the in-memory data provided.
307 Keys are EXTNAME values, which must be unique (no EXTVER disambiguation).
308 Precompressed pixel values are never copied or transferred to subsets.
309 """
311 def add_header(
312 self,
313 header: astropy.io.fits.Header,
314 name: str | None = None,
315 ver: int | None = None,
316 key: ExtensionKey | None = None,
317 ) -> None:
318 """Add a header to the opaque metadata if it is not already present,
319 and strip EXTNAME and EXTVER if present.
321 Parameters
322 ----------
323 header
324 Header to add. May be modified in place.
325 name
326 EXTNAME (all caps). If not provided, the EXTNAME card must be
327 present in the header. Use ``""`` for the primary header.
328 ver
329 EXTVER. If not provided and the EXTVER card is not present,
330 defaults to 1.
331 key
332 Combination of EXTNAME and EXTVER; used instead of both ``name``
333 and ``ver`` if provided.
334 """
335 if key is None:
336 if name is None:
337 name = header["EXTNAME"]
338 if ver is None:
339 ver = header.get("EXTVER", 1)
340 key = ExtensionKey(name, ver)
341 strip_butler_cards(header)
342 if key not in self.headers:
343 header.remove("EXTNAME", ignore_missing=True)
344 header.remove("EXTVER", ignore_missing=True)
345 self.headers[key] = header
347 def maybe_use_precompressed(self, name: str) -> astropy.io.fits.BinTableHDU | None:
348 """Look up the given EXTNAME to see if there is a tile compressed image
349 HDU that should be used directly, instead of requantizing.
351 Parameters
352 ----------
353 name
354 EXTNAME (all caps).
356 Returns
357 -------
358 `astropy.io.fits.BinTableHDU` | `None`
359 An already-compressed HDU, in binary table form, or `None` if there
360 is no precompressed HDU for this EXTNAME.
361 """
362 if (precompressed := self.precompressed.get(name)) is None:
363 return None
364 return astropy.io.fits.BinTableHDU(precompressed.data, header=precompressed.header.copy(), name=name)
366 def extract_legacy_primary_header(self, header: astropy.io.fits.Header) -> dict[str, Any]:
367 """Update the opaque metadata with the header of the primary HDU
368 of a legacy (`lsst.afw.image`) FITS file, stripping cards we know we
369 don't need and extracting any ``LSST IMAGES ...`` cards into a
370 dictionary we return.
371 """
372 primary_header = header.copy(strip=True)
373 # No idea what these spare TAN-SIP headers are doing in the afw
374 # FITS files, but we'll strip them here:
375 primary_header.remove("A_ORDER", ignore_missing=True)
376 primary_header.remove("B_ORDER", ignore_missing=True)
377 strip_legacy_exposure_cards(primary_header)
378 strip_butler_cards(primary_header)
379 metadata: dict[str, Any] = {}
380 for n in itertools.count():
381 if (key := header.pop(f"LSST IMAGES KEY {n + 1}", ...)) is ...:
382 break
383 value = header.pop(f"LSST IMAGES VALUE {n + 1}")
384 metadata[key] = value
385 self.headers[ExtensionKey()] = primary_header
386 return metadata
388 def copy(self) -> FitsOpaqueMetadata:
389 # Docstring inherited.
390 return FitsOpaqueMetadata(headers=self.headers)
392 def subset(self, bbox: Box) -> FitsOpaqueMetadata:
393 # Docstring inherited.
394 return FitsOpaqueMetadata(headers=self.headers)
396 def get_instrumental_unit(self) -> astropy.units.UnitBase | None:
397 """Extract the ``LSST ISR UNIT`` key from the primary header (if it
398 exists) and wrap it with Astropy.
399 """
400 if (primary_header := self.headers.get(ExtensionKey())) is not None:
401 if (instrumental_unit_str := primary_header.get("LSST ISR UNITS")) is not None:
402 return astropy.units.Unit(instrumental_unit_str)
403 return None
406def add_offset_wcs(header: astropy.io.fits.Header, *, x: int | float, y: int | float, key: str = "A") -> None:
407 """Add a trivial FITS WCS to a header that applies the appropriate offset
408 to map FITS array coordinates to a logical pixel grid.
410 Parameters
411 ----------
412 header
413 Header to update in-place.
414 x
415 Logical coordinate of the first column.
416 y
417 Logical coordinate of the first row.
418 key
419 Single-character suffix for this WCS.
420 """
421 header.set(f"CTYPE1{key}", "LINEAR")
422 header.set(f"CTYPE2{key}", "LINEAR")
423 header.set(f"CRPIX1{key}", 1.0)
424 header.set(f"CRPIX2{key}", 1.0)
425 header.set(f"CRVAL1{key}", float(x))
426 header.set(f"CRVAL2{key}", float(y))
427 header.set(f"CUNIT1{key}", "PIXEL")
428 header.set(f"CUNIT2{key}", "PIXEL")
431_WCS_VECTOR_KEYS = ("CUNIT", "CRPIX", "CRPIX", "CRVAL", "CRDELT", "CROTA", "CRDER", "CSYER", "CDELT")
432_WCS_MATRIX_KEYS = ("CD{0}_{1}", "PC{0}_{1}")
435def strip_wcs_cards(header: astropy.io.fits.Header) -> None:
436 """Strip WCS cards from a FITS header.
438 This does *not* attempt to cover all possible FITS WCS forms; it focuses on
439 the ones we actually plan to write (simple undistorted ones + TAN-SIP).
440 """
441 wcsaxes = header.pop("WCSAXES", 2)
442 for wcsname in [""] + list(string.ascii_uppercase):
443 header.remove("RADESYS" + wcsname, ignore_missing=True)
444 if "CTYPE1" + wcsname in header:
445 ctype: str = "" # just for linters that can't figure out that the loop always executes
446 for n in range(wcsaxes):
447 suffix = f"{n + 1}{wcsname}"
448 ctype = header.pop("CTYPE" + suffix)
449 for key in _WCS_VECTOR_KEYS:
450 header.remove(key + suffix, ignore_missing=True)
451 for m in range(wcsaxes):
452 for tmpl in _WCS_MATRIX_KEYS:
453 header.remove(tmpl.format(m + 1, suffix), ignore_missing=True)
454 if ctype.endswith("-SIP"):
455 _strip_sip_poly(header, wcsname, "A")
456 _strip_sip_poly(header, wcsname, "B")
457 _strip_sip_poly(header, wcsname, "AP")
458 _strip_sip_poly(header, wcsname, "BP")
459 header.remove("LONPOLE", ignore_missing=True)
460 header.remove("LATPOLE", ignore_missing=True)
461 header.remove("MJDREF", ignore_missing=True)
462 header.remove("PV1_3", ignore_missing=True)
463 header.remove("PV1_4", ignore_missing=True)
466def _strip_sip_poly(header: astropy.io.fits.Header, wcsname: str, which: str) -> None:
467 order: int | None = header.pop(f"{which}_ORDER{wcsname}", None)
468 if order is not None:
469 for i, j in itertools.product(range(order + 1), range(order + 1)):
470 header.remove(f"{which}_{i}_{j}{wcsname}", ignore_missing=True)
473def strip_legacy_exposure_cards(header: astropy.io.fits.Header) -> None:
474 """Strip header keywords added by lsst.afw.image.Exposure."""
475 header.remove("AR_HDU", ignore_missing=True)
476 for name in (
477 "FILTER",
478 "DETECTOR",
479 "VALID_POLYGON",
480 "SKYWCS",
481 "PSF",
482 "SUMMARYSTATS",
483 "AP_CORR_MAP",
484 "PHOTOCALIB",
485 ):
486 header.remove(f"{name}_ID", ignore_missing=True)
487 header.remove(f"ARCHIVE_ID_{name}", ignore_missing=True)
490def strip_butler_cards(header: astropy.io.fits.Header) -> None:
491 """Strip header keywords added by butler provenance that would be
492 incorrect if propagated to a downstream file.
493 """
494 for key in list(header):
495 if key.startswith("LSST BUTLER"):
496 del header[key]