Coverage for python / lsst / images / fits / _output_archive.py: 22%
201 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-03-05 02:32 -0800
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.
12from __future__ import annotations
14__all__ = ("FitsOutputArchive", "write")
16import dataclasses
17from collections import Counter
18from collections.abc import Callable, Hashable, Iterator, Mapping
19from contextlib import contextmanager
20from typing import Any, Self
22import astropy.io.fits
23import astropy.table
24import numpy as np
25import pydantic
27from .._transforms import FrameSet
28from ..serialization import (
29 ArchiveTree,
30 ArrayReferenceModel,
31 ColumnDefinitionModel,
32 NestedOutputArchive,
33 NumberType,
34 OutputArchive,
35 TableCellReferenceModel,
36 TableModel,
37 no_header_updates,
38)
39from ._common import ExtensionHDU, ExtensionKey, FitsCompressionOptions, FitsOpaqueMetadata
42def write(
43 obj: Any,
44 filename: str,
45 compression_options: Mapping[str, FitsCompressionOptions | None] | None = None,
46 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
47 compression_seed: int | None = None,
48) -> Any:
49 """Write an object with a ``serialize`` method to a FITS file.
51 Parameters
52 ----------
53 filename
54 Name of the file to write to. Must not already exist.
55 compression_options
56 Options for how to compress the FITS file, keyed by the name of
57 the attribute (with JSON pointer ``/`` separators for nested
58 attributes).
59 update_header
60 A callback that will be given the primary HDU FITS header and an
61 opportunity to modify it.
62 compression_seed
63 A FITS tile compression seed to use whenever the configured
64 compression seed is `None` or (for backwards compatibility) ``0``.
65 This value is then incremented every time it is used.
67 Returns
68 -------
69 `.serialization.ArchiveTree`
70 The serialized representation of the object.
71 """
72 opaque_metadata = getattr(obj, "_opaque_metadata", None)
73 name = getattr(obj, "_archive_default_name", None)
74 with FitsOutputArchive.open(
75 filename,
76 compression_options=compression_options,
77 opaque_metadata=opaque_metadata,
78 update_header=update_header,
79 compression_seed=compression_seed,
80 ) as archive:
81 tree = archive.serialize_direct(name, obj.serialize) if name is not None else obj.serialize(archive)
82 archive.add_tree(tree)
83 return tree
86class FitsOutputArchive(OutputArchive[TableCellReferenceModel]):
87 """An implementation of the `.serialization.OutputArchive` interface that
88 writes to FITS files.
90 Instances of this class should only be constructed via the `open`
91 context manager.
92 """
94 def __init__(
95 self,
96 hdu_list: astropy.io.fits.HDUList,
97 compression_options: Mapping[str, FitsCompressionOptions | None] | None = None,
98 opaque_metadata: Any = None,
99 compression_seed: int | None = None,
100 ):
101 # JSON blobs for objects we've saved as pointers:
102 self._pointer_targets: list[bytes] = []
103 # Mapping from user provided key (e.g. id(some object)) to a table
104 # pointer to where we actually saved it:
105 self._pointers_by_key: dict[Hashable, TableCellReferenceModel] = {}
106 self._hdu_list = hdu_list
107 self._primary_hdu = astropy.io.fits.PrimaryHDU()
108 # TODO: add subformat description and version to primary HDU.
109 self._primary_hdu.header.set("INDXADDR", 0, "Offset in bytes to the HDU index.")
110 self._primary_hdu.header.set("INDXSIZE", 0, "Size of the HDU index.")
111 self._primary_hdu.header.set("JSONADDR", 0, "Offset in bytes to the JSON tree HDU.")
112 self._primary_hdu.header.set("JSONSIZE", 0, "Size of the JSON tree HDU.")
113 self._hdus_by_name = Counter[str]()
114 self._compression_options = dict(compression_options) if compression_options is not None else {}
115 self._compression_seed = compression_seed
116 self._opaque_metadata = (
117 opaque_metadata if isinstance(opaque_metadata, FitsOpaqueMetadata) else FitsOpaqueMetadata()
118 )
119 if (opaque_primary_header := self._opaque_metadata.headers.get(ExtensionKey())) is not None:
120 self._primary_hdu.header.extend(opaque_primary_header)
121 self._hdu_list.append(self._primary_hdu)
122 self._json_hdu_added: bool = False
123 self._frame_sets: list[tuple[FrameSet, TableCellReferenceModel]] = []
125 @classmethod
126 @contextmanager
127 def open(
128 cls,
129 filename: str,
130 compression_options: Mapping[str, FitsCompressionOptions | None] | None = None,
131 opaque_metadata: Any = None,
132 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
133 compression_seed: int | None = None,
134 ) -> Iterator[Self]:
135 """Create an output archive that writes to the given file.
137 Parameters
138 ----------
139 filename
140 Name of the file to write to. Must not already exist.
141 compression_options
142 Options for how to compress the FITS file, keyed by the name of
143 the attribute (with JSON pointer ``/`` separators for nested
144 attributes).
145 opaque_metadata
146 Metadata read from an input archive along with the object being
147 written now. Ignored if the metadata is not from a FITS archive.
148 update_header
149 A callback that will be given the primary HDU FITS header and an
150 opportunity to modify it.
151 compression_seed
152 A FITS tile compression seed to use whenever the configured
153 compression seed is `None` or (for backwards compatibility) ``0``.
154 This value is then incremented every time it is used.
156 Returns
157 -------
158 `contextlib.AbstractContextManager` [`FitsOutputArchive`]
159 A context manager that returns a `FitsOutputArchive` when entered.
160 """
161 with astropy.io.fits.open(filename, mode="append") as hdu_list:
162 if hdu_list:
163 raise OSError(f"File {filename!r} already exists.")
164 archive = cls(hdu_list, compression_options, opaque_metadata, compression_seed=compression_seed)
165 update_header(hdu_list[0].header)
166 yield archive
167 if not archive._json_hdu_added:
168 raise RuntimeError("Write context exited without 'add_tree' being called.")
169 hdu_list.flush()
170 # This multi-open dance is necessary to get Astropy to tell us the
171 # byte addresses of the HDUs. Hopefully we can get an upstream change
172 # make this unnecessary at some point.
173 with astropy.io.fits.open(filename, mode="readonly", disable_image_compression=True) as hdu_list:
174 index_hdu = cls._make_index_table(hdu_list)
175 with astropy.io.fits.open(filename, mode="append") as hdu_list:
176 hdu_list.append(index_hdu)
177 json_bytes = _HDUBytes.from_index_row(index_hdu.data[-1])
178 index_bytes = _HDUBytes.from_write_hdu(index_hdu)
179 # Update the primary HDU with the address and size of the index and
180 # JSON HDUs, and rewrite just that. We do this write manually, since
181 # astropy's docs on its 'update' mode are scarce and it's not obvious
182 # whether we can guarantee it won't rewrite the whole file if we edit
183 # the primary header.
184 archive._primary_hdu.header["INDXADDR"] = index_bytes.header_address
185 archive._primary_hdu.header["INDXSIZE"] = index_bytes.size
186 archive._primary_hdu.header["JSONADDR"] = json_bytes.header_address
187 archive._primary_hdu.header["JSONSIZE"] = json_bytes.size
188 with open(filename, "r+b") as stream:
189 stream.write(archive._primary_hdu.header.tostring().encode())
191 def serialize_direct[T: pydantic.BaseModel](
192 self, name: str, serializer: Callable[[OutputArchive[TableCellReferenceModel]], T]
193 ) -> T:
194 nested = NestedOutputArchive[TableCellReferenceModel](name, self)
195 return serializer(nested)
197 def serialize_pointer[T: pydantic.BaseModel](
198 self, name: str, serializer: Callable[[OutputArchive[TableCellReferenceModel]], T], key: Hashable
199 ) -> TableCellReferenceModel:
200 if (pointer := self._pointers_by_key.get(key)) is not None:
201 return pointer
202 pointer = TableCellReferenceModel(
203 source="fits:JSON", column="JSON", row=len(self._pointer_targets) + 1
204 )
205 model = self.serialize_direct("", serializer)
206 self._pointer_targets.append(model.model_dump_json().encode())
207 self._pointers_by_key[key] = pointer
208 return pointer
210 def serialize_frame_set[T: pydantic.BaseModel](
211 self, name: str, frame_set: FrameSet, serializer: Callable[[OutputArchive], T], key: Hashable
212 ) -> TableCellReferenceModel:
213 # Docstring inherited.
214 pointer = self.serialize_pointer(name, serializer, key)
215 self._frame_sets.append((frame_set, pointer))
216 return pointer
218 def iter_frame_sets(self) -> Iterator[tuple[FrameSet, TableCellReferenceModel]]:
219 return iter(self._frame_sets)
221 def add_array(
222 self,
223 array: np.ndarray,
224 *,
225 name: str | None = None,
226 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
227 ) -> ArrayReferenceModel:
228 if name is None:
229 raise RuntimeError("Cannot save array with name=None unless it is nested.")
230 extname = name.upper()
231 hdu = self._opaque_metadata.maybe_use_precompressed(extname)
232 if hdu is None:
233 if (compression_options := self._get_compression_options(name)) is not None:
234 hdu = compression_options.make_hdu(array, name=extname)
235 else:
236 hdu = astropy.io.fits.ImageHDU(array, name=extname)
237 key = self._add_hdu(hdu, update_header)
238 return ArrayReferenceModel(
239 source=str(key),
240 shape=list(array.shape),
241 datatype=NumberType.from_numpy(array.dtype),
242 )
244 def add_table(
245 self,
246 table: astropy.table.Table,
247 *,
248 name: str | None = None,
249 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
250 ) -> TableModel:
251 if name is None:
252 raise RuntimeError("Cannot save table with name=None unless it is nested.")
253 extname = name.upper()
254 hdu: astropy.io.fits.BinTableHDU = astropy.io.fits.table_to_hdu(table, name=extname)
255 columns = ColumnDefinitionModel.from_record_dtype(hdu.data.dtype)
256 for c in columns:
257 c.update_from_table(table)
258 key = self._add_hdu(hdu, update_header)
259 return TableModel(source=str(key), columns=columns)
261 def add_structured_array(
262 self,
263 array: np.ndarray,
264 *,
265 name: str | None = None,
266 units: Mapping[str, astropy.units.Unit] | None = None,
267 descriptions: Mapping[str, str] | None = None,
268 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
269 ) -> TableModel:
270 if name is None:
271 raise RuntimeError("Cannot save structured array with name=None unless it is nested.")
272 extname = name.upper()
273 hdu = astropy.io.fits.BinTableHDU(array, name=extname)
274 columns = ColumnDefinitionModel.from_record_dtype(hdu.data.dtype)
275 if units is not None:
276 for c in columns:
277 c.unit = units.get(c.name)
278 if descriptions is not None:
279 for c in columns:
280 c.description = descriptions.get(c.name, "")
281 key = self._add_hdu(hdu, update_header)
282 return TableModel(source=str(key), columns=columns)
284 def _add_hdu(
285 self,
286 hdu: astropy.io.fits.ImageHDU | astropy.io.fits.CompImageHDU | astropy.io.fits.BinTableHDU,
287 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
288 ) -> ExtensionKey:
289 n_hdus = self._hdus_by_name.get(hdu.name, 0)
290 key = ExtensionKey(hdu.name, n_hdus + 1)
291 if n_hdus:
292 hdu.header["EXTVER"] = key.ver
293 self._hdus_by_name[hdu.name] += 1
294 update_header(hdu.header)
295 if (opaque_headers := self._opaque_metadata.headers.get(key)) is not None:
296 hdu.header.extend(opaque_headers)
297 self._hdu_list.append(hdu)
298 return key
300 def add_tree(self, tree: ArchiveTree) -> None:
301 """Write the JSON tree to the archive.
303 This method must be called exactly once, just before the `open` context
304 is exited.
306 Parameters
307 ----------
308 tree
309 Pydantic model that represents the tree.
310 """
311 json_hdu = astropy.io.fits.BinTableHDU.from_columns(
312 [astropy.io.fits.Column("JSON", "PB")],
313 nrows=len(self._pointer_targets) + 1,
314 name="JSON",
315 )
316 json_hdu.data[0]["JSON"] = np.frombuffer(tree.model_dump_json().encode(), dtype=np.byte)
317 for n, json_target_data in enumerate(self._pointer_targets):
318 json_hdu.data[n + 1]["JSON"] = np.frombuffer(json_target_data, dtype=np.byte)
319 self._hdu_list.append(json_hdu)
320 self._json_hdu_added = True
322 def _get_compression_options(self, name: str) -> FitsCompressionOptions | None:
323 result = self._compression_options.get(name, FitsCompressionOptions.DEFAULT)
324 if result is None or result.quantization is None:
325 return result
326 if self._compression_seed is not None and not result.quantization.seed:
327 result = result.model_copy(
328 update={
329 "quantization": result.quantization.model_copy(update={"seed": self._compression_seed})
330 }
331 )
332 self._compression_seed += 1
333 if self._compression_seed > 10000:
334 self._compression_seed = 1
335 # MyPy can tell that result.quantization is not None in the 'if', but
336 # forgets that by this 'else':
337 elif result.quantization.seed is None: # type: ignore[union-attr]
338 raise RuntimeError("No quantization seed provided.")
339 return result
341 @staticmethod
342 def _make_index_table(hdu_list: astropy.io.fits.HDUList) -> astropy.io.fits.BinTableHDU:
343 # We use a fixed-length string for the EXTNAME column; it might be
344 # better to use a variable-length array, but I have not been able to
345 # figure out how to get astropy to accept a string for the the
346 # character (TFORM='A') variant of that. And that's only better if the
347 # EXTNAMEs get super long, which is not likely (but maybe something to
348 # guard against).
349 max_name_size = max(len(hdu.header.get("EXTNAME", "")) for hdu in hdu_list)
350 index_hdu = astropy.io.fits.BinTableHDU.from_columns(
351 [
352 astropy.io.fits.Column("EXTNAME", f"A{max_name_size}"),
353 astropy.io.fits.Column("EXTVER", "J"),
354 astropy.io.fits.Column("XTENSION", "A8"),
355 astropy.io.fits.Column("ZIMAGE", "L"),
356 ]
357 + _HDUBytes.get_index_hdu_columns(),
358 nrows=len(hdu_list),
359 name="INDEX",
360 )
361 hdu: ExtensionHDU | astropy.io.fits.PrimaryHDU
362 for n, hdu in enumerate(hdu_list):
363 index_hdu.data[n]["EXTNAME"] = hdu.header.get("EXTNAME", "")
364 index_hdu.data[n]["EXTVER"] = hdu.header.get("EXTVER", 1)
365 index_hdu.data[n]["XTENSION"] = hdu.header.get("XTENSION", "IMAGE")
366 index_hdu.data[n]["ZIMAGE"] = hdu.header.get("ZIMAGE", False)
367 bytes = _HDUBytes.from_read_hdu(hdu)
368 bytes.update_index_row(index_hdu.data[n])
369 return index_hdu
372@dataclasses.dataclass
373class _HDUBytes:
374 """A struct that records the byte offsets into a FITS HDU."""
376 @classmethod
377 def from_write_hdu(cls, hdu: astropy.io.fits.PrimaryHDU | ExtensionHDU) -> Self:
378 """Construct from an Astropy HDU instance that has just been written.
380 Parameters
381 ----------
382 hdu
383 An Astropy HDU object.
385 Returns
386 -------
387 hdu_bytes
388 Struct with byte offsets.
390 Notes
391 -----
392 This method relies on internal Astropy attributes and does not work on
393 CompImageHDU objects.
394 """
395 # This is implemented by accessing private Astropy attributes because
396 # it turns out that's much more reliable than the public fileinfo()
397 # method, which seems to always return a dict with `None` entries or
398 # raise; it looks buggy, but docs are scarce enough that it's not clear
399 # what the right behavior is supposed to be.
400 if (header_address := getattr(hdu, "_header_offset", None)) is None:
401 raise RuntimeError("Failed to get Astropy's _header_offset.")
402 if (data_address := getattr(hdu, "_data_offset", None)) is None:
403 raise RuntimeError("Failed to get Astropy's _data_offset.")
404 if (data_size := getattr(hdu, "_data_size", None)) is None:
405 raise RuntimeError("Failed to get Astropy's _data_size.")
406 return cls(header_address, data_address, data_size)
408 @classmethod
409 def from_read_hdu(cls, hdu: astropy.io.fits.PrimaryHDU | ExtensionHDU) -> Self:
410 """Construct from an Astropy HDU instance that has just been read.
412 Parameters
413 ----------
414 hdu
415 An Astropy HDU object.
417 Returns
418 -------
419 hdu_bytes
420 Struct with byte offsets.
421 """
422 info = hdu.fileinfo()
423 header_address = info["hdrLoc"]
424 data_address = info["datLoc"]
425 data_size = info["datSpan"]
426 return cls(header_address, data_address, data_size)
428 @classmethod
429 def from_index_row(cls, row: np.void) -> Self:
430 """Construct from a row of the index HDU.
432 Parameters
433 ----------
434 row
435 A Numpy struct-like scalar.
437 Returns
438 -------
439 hdu_bytes
440 Struct with byte offsets.
441 """
442 return cls(
443 header_address=int(row["HDRADDR"]),
444 data_address=int(row["DATADDR"]),
445 data_size=int(row["DATSIZE"]),
446 )
448 @staticmethod
449 def get_index_hdu_columns() -> list[astropy.io.fits.Column]:
450 """Return the definitions of the columns this class gets and sets
451 from the index HDU.
453 Returns
454 -------
455 columns
456 A `list` of `astropy.io.fits.Column` objects that represent the
457 header address, data address, and data size.
458 """
459 return [
460 astropy.io.fits.Column("HDRADDR", "K"),
461 astropy.io.fits.Column("DATADDR", "K"),
462 astropy.io.fits.Column("DATSIZE", "K"),
463 ]
465 header_address: int
466 """Offset from the beginning of the start of the file to the header of this
467 HDU, in bytes.
468 """
470 data_address: int
471 """Offset from the beginning of the start of the file to the data section
472 of this HDU, in bytes.
473 """
475 data_size: int
476 """Size of the data section in bytes."""
478 @property
479 def header_size(self) -> int:
480 """Size of the header in bytes."""
481 return self.data_address - self.header_address
483 @property
484 def end_address(self) -> int:
485 """Offset in bytes from the start of the file to the end of the HDU."""
486 return self.data_address + self.data_size
488 @property
489 def size(self) -> int:
490 """Total size of this HDU in bytes."""
491 return self.data_size + self.data_address - self.header_address
493 def update_index_row(self, row: np.void) -> None:
494 """Set the values of a row of the index HDU from this strut.
496 Parameters
497 ----------
498 row
499 A Numpy struct-like scalar to modify in place.
500 """
501 row["HDRADDR"] = self.header_address
502 row["DATADDR"] = self.data_address
503 row["DATSIZE"] = self.data_size