Coverage for python/lsst/images/ndf/_output_archive.py: 80%
350 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-23 09:23 +0000
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-23 09:23 +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.
12from __future__ import annotations
14__all__ = (
15 "NdfOutputArchive",
16 "write",
17)
19import os
20from collections.abc import Callable, Hashable, Iterator, Mapping, Sequence
21from contextlib import contextmanager
22from pathlib import Path
23from typing import Any, Self
25import astropy.io.fits
26import astropy.table
27import astropy.units
28import h5py
29import numpy as np
30import pydantic
32from .._color_image import ColorImage
33from .._transforms import FrameSet
34from .._transforms._ast import Channel, CmpFrame, CmpMap, ShiftMap, StringStream, UnitMap
35from .._transforms._ast import Frame as AstFrame
36from .._transforms._ast import FrameSet as AstFrameSet
37from .._transforms._transform import _prepend_ast_shift
38from ..fits._common import ExtensionKey, FitsOpaqueMetadata
39from ..serialization import (
40 ArchiveTree,
41 ArrayReferenceModel,
42 ButlerInfo,
43 MetadataValue,
44 NestedOutputArchive,
45 NumberType,
46 OutputArchive,
47 TableColumnModel,
48 TableModel,
49 no_header_updates,
50)
51from . import _hds
52from ._common import (
53 HdsNameShrinker,
54 NdfPointerModel,
55 archive_path_to_hdf5_path,
56 archive_path_to_hdf5_path_components,
57)
58from ._model import HdsPrimitive, HdsStructure, Ndf, NdfArray, NdfContainer, NdfDocument, NdfQuality, NdfWcs
60_NDF_FORMAT_VERSION = 1
61"""Container layout version for files written by `NdfOutputArchive`.
63Bumps when the NDF layout (NdfDocument shape, .MORE.LSST contents)
64changes. Independent of any data-model ``SCHEMA_VERSION``.
65"""
68def write(
69 obj: Any,
70 filename: str | Path | None = None,
71 *,
72 metadata: dict[str, MetadataValue] | None = None,
73 butler_info: ButlerInfo | None = None,
74 compression_options: Mapping[str, Any] | None = None,
75) -> ArchiveTree:
76 """Write a serializable object to an NDF (HDS-on-HDF5) file.
78 Parameters
79 ----------
80 obj
81 Object with a ``serialize`` method. May carry an
82 ``_opaque_metadata`` attribute (a
83 `~lsst.images.fits.FitsOpaqueMetadata`)
84 whose primary-HDU header gets written to ``/MORE/FITS``. This
85 preserves FITS cards on objects that originated from a FITS read;
86 butler provenance is conveyed through ``butler_info`` instead.
87 filename
88 Path to write to. Must not already exist. If `None`, an
89 in-memory HDF5 file is used and the on-disk artefact is
90 discarded; the returned tree still reflects all the writes the
91 archive made (useful for tests).
92 metadata, butler_info
93 Optional caller-supplied entries that are written into the
94 returned `~lsst.images.serialization.ArchiveTree`.
95 compression_options
96 Optional dict forwarded to the archive constructor for h5py
97 dataset compression.
99 Returns
100 -------
101 `~lsst.images.serialization.ArchiveTree`
102 The Pydantic tree the object's ``serialize`` produced (with
103 ``metadata``/``butler_info`` applied).
104 """
105 opaque_metadata = getattr(obj, "_opaque_metadata", None)
106 if not isinstance(opaque_metadata, FitsOpaqueMetadata):
107 opaque_metadata = FitsOpaqueMetadata()
108 archive_default_name = getattr(obj, "_archive_default_name", None)
109 with NdfOutputArchive.open(
110 filename,
111 compression_options=compression_options,
112 opaque_metadata=opaque_metadata,
113 **_get_archive_layout(obj),
114 ) as archive:
115 tree = archive.serialize_root(obj, metadata, butler_info)
116 archive.add_tree(
117 tree,
118 sky_projection=getattr(obj, "sky_projection", None),
119 bbox=getattr(obj, "bbox", None),
120 unit=getattr(obj, "unit", None),
121 root_name=archive_default_name or type(obj).__name__,
122 )
123 return tree
126def _origin_from_bbox(bbox: Any) -> tuple[int, ...]:
127 """Extract NDF/Fortran-order origin tuple from an lsst.images Box.
129 The exact attribute names on Box depend on the version. Inspect the
130 object and pick whatever exposes the integer lower bounds. For a 2D
131 image with bbox lower bound (x_min, y_min) the returned tuple is
132 ``(x_min, y_min)``.
133 """
134 # Box exposes .x and .y properties returning Interval objects with a
135 # .start attribute (the lower bound).
136 if hasattr(bbox, "x") and hasattr(bbox, "y"): 136 ↛ 141line 136 didn't jump to line 141 because the condition on line 136 was always true
137 x = bbox.x
138 y = bbox.y
139 if hasattr(x, "start") and hasattr(y, "start"): 139 ↛ 141line 139 didn't jump to line 141 because the condition on line 139 was always true
140 return (int(x.start), int(y.start))
141 raise AttributeError(
142 f"Don't know how to extract origin from bbox of type {type(bbox).__name__!r}; "
143 "_origin_from_bbox needs updating."
144 )
147def _unit_to_ndf_string(unit: astropy.units.UnitBase) -> str:
148 """Return an ASCII unit string for the NDF UNITS component."""
149 try:
150 return unit.to_string(format="fits")
151 except ValueError:
152 return unit.to_string()
155def _fits_header_records(header: astropy.io.fits.Header) -> list[str]:
156 """Return fixed-width FITS records for an opaque NDF FITS extension.
158 NDF ``.MORE.FITS`` is a ``_CHAR*80`` vector. Use Astropy's FITS
159 header serializer to preserve multi-record logical cards such as
160 CONTINUE strings, but omit the FITS END card and 2880-byte padding
161 because this is an NDF extension component, not a complete FITS
162 header block.
163 """
164 block = header.tostring(sep="", endcard=False, padding=False)
165 encoded = block.encode("ascii")
166 if len(encoded) % 80: 166 ↛ 167line 166 didn't jump to line 167 because the condition on line 166 was never true
167 raise ValueError(
168 f"FITS header block is {len(encoded)} bytes, not a multiple of the 80-byte FITS record size."
169 )
170 return [block[n : n + 80] for n in range(0, len(block), 80)]
173def _get_archive_layout(obj: Any) -> dict[str, Any]:
174 """Return NDF document layout options for a top-level object."""
175 if isinstance(obj, ColorImage):
176 return {
177 "root": NdfContainer(),
178 "lsst_path": "/LSST",
179 "direct_ndf_array_paths": {
180 "red": "/RED",
181 "green": "/GREEN",
182 "blue": "/BLUE",
183 },
184 "wcs_ndf_paths": ("/RED", "/GREEN", "/BLUE"),
185 }
186 return {}
189def _show_ast_for_ndf(ast_frame_set: Any, bbox: Any | None) -> str:
190 """Return AST Channel text matching Starlink NDF WCS serialization.
192 Tags the original base frame with ``Domain="PIXEL"`` and prepends a
193 new ``Domain="GRID"`` base frame related to it by a `ShiftMap` whose
194 shift converts ``bbox``-origin pixel coordinates into 1-based grid
195 coordinates. The result is written via an abstraction-layer
196 ``Channel`` configured with the same options the Starlink C writer
197 uses (``Full=-1,Comment=0``; see ``ndf1Wwrt.c``) plus ``Indent=0`` so
198 each line is just the bare AST item with the single-space prefix
199 that ``ndf1Rdast`` strips back off on read.
200 """
201 if bbox is None: 201 ↛ 202line 201 didn't jump to line 202 because the condition on line 201 was never true
202 x_shift = 1.0
203 y_shift = 1.0
204 else:
205 x_shift = 1.0 - float(bbox.x.start)
206 y_shift = 1.0 - float(bbox.y.start)
208 saved_current = ast_frame_set.current
209 ast_frame_set.current = ast_frame_set.base
210 ast_frame_set.domain = "PIXEL"
211 ast_frame_set.current = saved_current
212 _prepend_ast_shift(ast_frame_set, x=x_shift, y=y_shift, ast_domain="GRID")
214 stream = StringStream()
215 channel = Channel(stream, options="Full=-1,Comment=0,Indent=0")
216 channel.write(ast_frame_set)
217 return stream.getSinkData()
220def _show_mask_ast_for_ndf(
221 parent_ast_frame_set: Any,
222 origin: Sequence[int],
223 *,
224 labels: Sequence[str] = (),
225) -> str:
226 """Return an NDF WCS for the 3D native mask sub-NDF.
228 The first two axes reuse the parent image's pixel-to-sky mapping. The
229 third axis is a generic mask-byte coordinate that passes through unchanged.
230 """
231 n_axes = len(origin)
232 ast_frame_set = AstFrameSet(AstFrame(n_axes, "Domain=GRID"))
233 pixel_frame = AstFrame(n_axes, "Domain=PIXEL")
234 for axis, label in enumerate(labels[:n_axes], start=1):
235 pixel_frame.setLabel(axis, label)
236 shifts = [1.0 - float(axis_origin) for axis_origin in origin]
237 ast_frame_set.addFrame(AstFrameSet.BASE, ShiftMap(shifts), pixel_frame)
239 parent_pixel_to_sky = parent_ast_frame_set.getMapping(
240 parent_ast_frame_set.base,
241 parent_ast_frame_set.current,
242 )
243 parent_sky_frame = parent_ast_frame_set.getFrame(parent_ast_frame_set.current)
244 mask_axis_frame = AstFrame(1, "Domain=MASK")
245 mask_axis_frame.setLabel(1, labels[2] if len(labels) > 2 else "mask-byte")
246 ast_frame_set.addFrame(
247 ast_frame_set.current,
248 CmpMap(parent_pixel_to_sky, UnitMap(1), False),
249 CmpFrame(parent_sky_frame, mask_axis_frame),
250 )
252 stream = StringStream()
253 channel = Channel(stream, options="Full=-1,Comment=0,Indent=0")
254 channel.write(ast_frame_set)
255 return stream.getSinkData()
258class NdfOutputArchive(OutputArchive[NdfPointerModel]):
259 """An `~lsst.images.serialization.OutputArchive` implementation
260 that writes HDS-on-HDF5 files compatible with the Starlink NDF data
261 model.
263 Parameters
264 ----------
265 file
266 An open `h5py.File` opened in a writable mode. The archive does
267 not close the file; the caller is responsible for that.
268 compression_options
269 Optional dict passed through to `h5py.Group.create_dataset` for image
270 arrays (e.g. ``{"compression": "gzip", "compression_opts": 4}``).
271 opaque_metadata
272 Optional `~lsst.images.fits.FitsOpaqueMetadata`; if its primary-HDU
273 header is non-empty its cards will be written to ``/MORE/FITS`` by the
274 top-level `write` function.
275 root
276 Root NDF or NDF container to populate; a new empty ``Ndf`` is created
277 when `None`.
278 lsst_path
279 HDS path under which the LSST extension is written.
280 direct_ndf_array_paths
281 Mapping from archive path to NDF array path for arrays written
282 directly as NDF components rather than into the LSST extension.
283 wcs_ndf_paths
284 NDF paths for which WCS information is written.
285 """
287 def __init__(
288 self,
289 file: h5py.File,
290 compression_options: Mapping[str, Any] | None = None,
291 opaque_metadata: FitsOpaqueMetadata | None = None,
292 root: Ndf | NdfContainer | None = None,
293 lsst_path: str = "/MORE/LSST",
294 direct_ndf_array_paths: Mapping[str, str] | None = None,
295 wcs_ndf_paths: Sequence[str] = ("/",),
296 ) -> None:
297 super().__init__()
298 self._file = file
299 self._document = NdfDocument(root=root if root is not None else Ndf())
300 self._lsst_path = lsst_path.rstrip("/") or "/LSST"
301 self._direct_ndf_array_paths = dict(direct_ndf_array_paths) if direct_ndf_array_paths else {}
302 self._wcs_ndf_paths = tuple(wcs_ndf_paths)
303 self._bbox_array_struct_paths: set[str] = set()
304 self._compression_options = dict(compression_options) if compression_options else {}
305 self._opaque_metadata = opaque_metadata if opaque_metadata is not None else FitsOpaqueMetadata()
306 self._frame_sets: list[tuple[FrameSet, NdfPointerModel]] = []
307 self._pointers: dict[Hashable, NdfPointerModel] = {}
308 self._hdf5_path_owners: dict[str, str] = {}
309 self._name_shrinker = HdsNameShrinker()
310 # Keep the open file in sync so existing direct-archive tests can
311 # inspect it immediately, while all mutations go through the IR.
312 self._flush()
314 @classmethod
315 @contextmanager
316 def open(
317 cls,
318 filename: str | Path | None,
319 *,
320 compression_options: Mapping[str, Any] | None = None,
321 opaque_metadata: FitsOpaqueMetadata | None = None,
322 root: Ndf | NdfContainer | None = None,
323 lsst_path: str = "/MORE/LSST",
324 direct_ndf_array_paths: Mapping[str, str] | None = None,
325 wcs_ndf_paths: Sequence[str] = ("/",),
326 ) -> Iterator[Self]:
327 """Open an NDF file for writing and yield an `NdfOutputArchive`.
329 ``filename=None`` uses an in-memory HDF5 file; the on-disk
330 artefact is discarded but the archive's writes still produce a
331 usable returned tree (handy for tests).
333 Parameters
334 ----------
335 filename
336 Name of the file to write to, or `None` to use an in-memory
337 HDF5 file whose on-disk artefact is discarded.
338 compression_options
339 Optional dict passed through to `h5py.Group.create_dataset`
340 for image arrays.
341 opaque_metadata
342 Optional `~lsst.images.fits.FitsOpaqueMetadata` whose cards are
343 written to ``/MORE/FITS`` by the top-level `write` function.
344 root
345 Root NDF or NDF container to populate; a new empty ``Ndf`` is
346 created when `None`.
347 lsst_path
348 HDS path under which the LSST extension is written.
349 direct_ndf_array_paths
350 Mapping from archive path to NDF array path for arrays written
351 directly as NDF components rather than into the LSST extension.
352 wcs_ndf_paths
353 NDF paths for which WCS information is written.
354 """
355 if filename is None: 355 ↛ 356line 355 didn't jump to line 356 because the condition on line 355 was never true
356 h5_file = h5py.File("inmem.sdf", "w", driver="core", backing_store=False)
357 else:
358 if os.path.exists(filename) and os.path.getsize(filename) > 0: 358 ↛ 359line 358 didn't jump to line 359 because the condition on line 358 was never true
359 raise OSError(f"File {filename!r} already exists.")
360 h5_file = h5py.File(filename, "w")
361 try:
362 yield cls(
363 h5_file,
364 compression_options=compression_options,
365 opaque_metadata=opaque_metadata,
366 root=root,
367 lsst_path=lsst_path,
368 direct_ndf_array_paths=direct_ndf_array_paths,
369 wcs_ndf_paths=wcs_ndf_paths,
370 )
371 finally:
372 h5_file.close()
374 def add_tree(
375 self,
376 tree: ArchiveTree,
377 *,
378 sky_projection: Any = None,
379 bbox: Any = None,
380 unit: astropy.units.UnitBase | None = None,
381 root_name: str | None = None,
382 ) -> None:
383 """Finalize the file: write WCS, units, JSON tree, and ORIGIN.
385 Writes the canonical NDF ``/WCS`` HDS structure (an AST channel
386 text dump that KAPPA / hdstrace expect) when ``sky_projection`` is
387 provided. A native mask sub-NDF at ``/MORE/LSST/MASK`` gets a 3D
388 WCS whose first two axes reuse the parent's sky projection and
389 whose third axis is the mask-byte coordinate. The JSON tree at
390 ``<lsst_path>/JSON`` remains the source of truth for symmetric
391 round-trips; ``/WCS`` is for Starlink tools. Auto-detect read of
392 ``/WCS/DATA`` into a typed ``SkyProjection`` is a follow-up.
394 Parameters
395 ----------
396 tree
397 Pydantic tree returned by the object's ``serialize`` method,
398 with ``metadata``/``butler_info`` already applied.
399 sky_projection, bbox, unit
400 Top-level object attributes that drive NDF-canonical writes.
401 root_name
402 Value to assign to the root group's ``HDS_ROOT_NAME``
403 attribute (fixed-length ASCII so KAPPA / hdstrace decode it).
404 """
405 if sky_projection is not None:
406 self._write_wcs(sky_projection, bbox)
407 if unit is not None and isinstance(self._document.root, Ndf):
408 self._document.ensure_ndf("/").set_units(_unit_to_ndf_string(unit))
409 json_text = tree.model_dump_json()
410 # Let ensure_structure derive types from path-component names so
411 # /MORE/LSST/... gets <MORE> stay as EXT and <LSST> typed LSST,
412 # mirroring HDS convention.
413 lsst = self._ensure_model_structure(self._lsst_path)
414 lsst.children["JSON"] = HdsPrimitive.char_array([json_text], width=max(80, len(json_text)))
415 lsst.children["DATA_MODEL"] = HdsPrimitive.char_scalar(
416 tree.schema_url, width=max(80, len(tree.schema_url))
417 )
418 lsst.children["FORMAT_VERSION"] = HdsPrimitive.array(np.array(_NDF_FORMAT_VERSION, dtype=np.int32))
419 primary = self._opaque_metadata.headers.get(ExtensionKey())
420 if primary is not None and len(primary):
421 cards = _fits_header_records(primary)
422 more = self._ensure_model_structure("/MORE")
423 more.children["FITS"] = HdsPrimitive.char_array(cards, width=80)
424 if bbox is not None: 424 ↛ 429line 424 didn't jump to line 429 because the condition on line 424 was always true
425 origin = _origin_from_bbox(bbox)
426 for struct_path in self._bbox_array_struct_paths:
427 if self._has_model_path(struct_path): 427 ↛ 426line 427 didn't jump to line 426 because the condition on line 427 was always true
428 self.set_array_origin(struct_path, origin)
429 if root_name is not None: 429 ↛ 431line 429 didn't jump to line 431 because the condition on line 429 was always true
430 self._document.root_name = root_name
431 self._flush()
433 def _write_wcs(self, sky_projection: Any, bbox: Any) -> None:
434 ast_frame_set = sky_projection._pixel_to_sky._get_ast_frame_set()
435 text = _show_ast_for_ndf(ast_frame_set, bbox)
436 lines = _hds.encode_ndf_ast_data(text)
437 for ndf_path in self._wcs_ndf_paths:
438 if self._has_model_path(ndf_path): 438 ↛ 437line 438 didn't jump to line 437 because the condition on line 438 was always true
439 self._document.ensure_ndf(ndf_path).set_wcs(NdfWcs(lines))
440 if self._has_model_path("/MORE/LSST/MASK"):
441 mask_origin = _origin_from_bbox(bbox) if bbox is not None else (0, 0)
442 mask_ndim = self._model_array_ndim("/MORE/LSST/MASK/DATA_ARRAY")
443 mask_origin = (*mask_origin, *((0,) * max(0, mask_ndim - len(mask_origin))))
444 mask_ast_frame_set = sky_projection._pixel_to_sky._get_ast_frame_set()
445 mask_text = _show_mask_ast_for_ndf(
446 mask_ast_frame_set,
447 mask_origin,
448 labels=("x", "y", "mask-byte"),
449 )
450 self._document.ensure_ndf("/MORE/LSST/MASK").set_wcs(NdfWcs(_hds.encode_ndf_ast_data(mask_text)))
452 def serialize_direct[T: pydantic.BaseModel | None](
453 self, name: str, serializer: Callable[[OutputArchive[NdfPointerModel]], T]
454 ) -> T:
455 nested = NestedOutputArchive[NdfPointerModel](name, self)
456 return serializer(nested)
458 def serialize_pointer[T: ArchiveTree](
459 self, name: str, serializer: Callable[[OutputArchive[NdfPointerModel]], T], key: Hashable
460 ) -> NdfPointerModel:
461 if (pointer := self._pointers.get(key)) is not None:
462 return pointer
463 archive_path = name if name.startswith("/") else f"/{name}"
464 target_path = self._archive_path_to_hdf5_path(archive_path)
465 self._register_hdf5_path(target_path, archive_path)
466 # Run the serializer first so any nested add_array / serialize_pointer
467 # calls write into the file before we dump this sub-tree to JSON.
468 model = self.serialize_direct(name, serializer)
469 json_text = model.model_dump_json()
470 # Store the JSON document as a "JSON" child of the target structure,
471 # mirroring the main /MORE/LSST/JSON tree. Writing it at the target
472 # path itself would clobber any nested arrays the serializer added
473 # under that path (IR.write_to_hdf5 deletes the existing node before
474 # writing a primitive).
475 # ensure_structure derives the HDS type from the leaf name, so
476 # /MORE/LSST/PSF is typed <PSF> rather than the generic <EXT>.
477 target = self._ensure_model_structure(target_path)
478 target.children["JSON"] = HdsPrimitive.char_array([json_text], width=max(80, len(json_text)))
479 self._flush()
480 pointer = NdfPointerModel(path=f"{target_path}/JSON")
481 self._pointers[key] = pointer
482 return pointer
484 def serialize_frame_set[T: ArchiveTree](
485 self, name: str, frame_set: FrameSet, serializer: Callable[[OutputArchive], T], key: Hashable
486 ) -> NdfPointerModel:
487 pointer = self.serialize_pointer(name, serializer, key)
488 self._frame_sets.append((frame_set, pointer))
489 return pointer
491 def iter_frame_sets(self) -> Iterator[tuple[FrameSet, NdfPointerModel]]:
492 return iter(self._frame_sets)
494 _COMPATIBLE_MASK_DTYPES = (np.dtype(np.uint8),)
495 _prefer_native_mask_arrays = True
497 def add_array(
498 self,
499 array: np.ndarray,
500 *,
501 name: str | None = None,
502 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
503 tile_shape: tuple[int, ...] | None = None,
504 options_name: str | None = None,
505 ) -> ArrayReferenceModel:
506 if name is None: 506 ↛ 507line 506 didn't jump to line 507 because the condition on line 506 was never true
507 raise ValueError("Anonymous arrays are not supported in the NDF archive.")
508 name, version = self._register_name(name)
509 # Recognised top-level names go to standard NDF locations.
510 # Anything else hoists under /MORE/LSST.
511 if name == "image":
512 root = self._document.ensure_ndf("/")
513 root.set_array_component(
514 "DATA_ARRAY",
515 array,
516 origin=np.zeros(array.ndim, dtype=np.int64),
517 compression_options=self._compression_options,
518 )
519 path = "/DATA_ARRAY/DATA"
520 self._bbox_array_struct_paths.add("/DATA_ARRAY")
521 elif name == "variance":
522 root = self._document.ensure_ndf("/")
523 root.set_array_component(
524 "VARIANCE",
525 array,
526 origin=np.zeros(array.ndim, dtype=np.int64),
527 compression_options=self._compression_options,
528 )
529 path = "/VARIANCE/DATA"
530 self._bbox_array_struct_paths.add("/VARIANCE")
531 elif name == "mask":
532 if array.ndim == 2 and array.dtype in self._COMPATIBLE_MASK_DTYPES:
533 self._set_quality_array(array)
534 path = "/QUALITY/QUALITY/DATA"
535 self._bbox_array_struct_paths.add("/QUALITY/QUALITY")
536 else:
537 # Native Mask serialization passes HDF5 shape
538 # (mask-byte, y, x). HDS reports the reverse dimension order,
539 # so Starlink tools see (x, y, mask-byte).
540 if array.ndim == 3 and array.dtype in self._COMPATIBLE_MASK_DTYPES: 540 ↛ 543line 540 didn't jump to line 543 because the condition on line 540 was always true
541 self._set_quality_array(self._collapse_mask_to_quality(array))
542 self._bbox_array_struct_paths.add("/QUALITY/QUALITY")
543 mask_ndf = self._document.ensure_ndf("/MORE/LSST/MASK")
544 mask_ndf.set_array_component(
545 "DATA_ARRAY",
546 array,
547 origin=np.zeros(array.ndim, dtype=np.int64),
548 bad_pixel=False,
549 compression_options=self._compression_options,
550 )
551 path = "/MORE/LSST/MASK/DATA_ARRAY/DATA"
552 self._bbox_array_struct_paths.add("/MORE/LSST/MASK/DATA_ARRAY")
553 elif name in self._direct_ndf_array_paths:
554 sub_ndf_path = self._direct_ndf_array_paths[name]
555 sub_ndf = self._document.ensure_ndf(sub_ndf_path)
556 sub_ndf.set_array_component(
557 "DATA_ARRAY",
558 array,
559 origin=np.zeros(array.ndim, dtype=np.int64),
560 compression_options=self._compression_options,
561 )
562 path = f"{sub_ndf_path}/DATA_ARRAY/DATA"
563 self._bbox_array_struct_paths.add(f"{sub_ndf_path}/DATA_ARRAY")
564 else:
565 # Hoisted numeric arrays are wrapped as sub-NDFs under
566 # /MORE/LSST/<UPPER_PATH> so Starlink tools (KAPPA `display`,
567 # `hdstrace`, etc.) can inspect them just like the main image.
568 # The sub-NDF has the canonical layout: top-level group with
569 # CLASS="NDF" containing a DATA_ARRAY structure (CLASS="ARRAY")
570 # with DATA + ORIGIN primitives. Over-long components are shrunk
571 # to fit the HDS object-name limit (DAT__SZNAM); repeated names
572 # get a version suffix on their leaf so siblings stay distinct.
573 archive_path, logical_id = self._versioned_archive_path(name, version)
574 sub_ndf_path = self._archive_path_to_hdf5_path(archive_path)
575 self._register_hdf5_path(sub_ndf_path, logical_id)
576 sub_ndf = self._document.ensure_ndf(sub_ndf_path)
577 sub_ndf.set_array_component(
578 "DATA_ARRAY",
579 array,
580 origin=np.zeros(array.ndim, dtype=np.int64),
581 compression_options=self._compression_options,
582 )
583 path = f"{sub_ndf_path}/DATA_ARRAY/DATA"
584 self._flush()
585 # Shape is stored in the JSON tree (matching the FITS archive) because
586 # MaskSerializationModel.bbox needs it before any arrays are read.
587 # Future work: resolve shape from the HDF5 dataset on read instead.
588 return ArrayReferenceModel(
589 source=f"ndf:{path}",
590 shape=list(array.shape),
591 datatype=NumberType.from_numpy(array.dtype),
592 )
594 def _ensure_path(self, path: str) -> h5py.Group:
595 """Walk/create groups for an HDF5 absolute path.
597 Intermediate groups created on demand are tagged with
598 ``CLASS="EXT"``, the HDS type for general-purpose extension
599 containers (matches what hds-v5 writes for ``MORE``).
600 """
601 self._ensure_model_structure(path, "EXT")
602 self._flush()
603 return self._file["/" if path in ("", "/") else path]
605 def _ensure_struct(self, path: str, hds_type: str) -> h5py.Group:
606 """Ensure a structure exists at ``path`` with the given HDS type."""
607 self._ensure_model_structure(path, hds_type)
608 self._flush()
609 return self._file["/" if path in ("", "/") else path]
611 def _ensure_array_structure(self, path: str) -> h5py.Group:
612 """Ensure an HDS ``ARRAY`` structure exists at ``path``."""
613 return self._ensure_struct(path, "ARRAY")
615 def _ensure_quality_structure(self) -> h5py.Group:
616 """Ensure ``/QUALITY`` exists with ``CLASS="QUALITY"`` and BADBITS.
618 BADBITS is set to 255 so every bit of the collapsed single-byte
619 QUALITY plane is treated as bad by NDF applications.
620 """
621 root = self._document.ensure_ndf("/")
622 if "QUALITY" not in root.children:
623 root.children["QUALITY"] = HdsStructure("QUALITY")
624 quality = root.get_structure("QUALITY")
625 quality.hds_type = "QUALITY"
626 quality.children["BADBITS"] = HdsPrimitive.array(np.array(255, dtype=np.uint8))
627 self._flush()
628 return self._file["/QUALITY"]
630 def _ensure_quality_array_structure(self) -> h5py.Group:
631 """Ensure the nested ``/QUALITY/QUALITY`` ARRAY structure exists."""
632 root = self._document.ensure_ndf("/")
633 if "QUALITY" not in root.children:
634 root.children["QUALITY"] = HdsStructure("QUALITY")
635 quality = root.get_structure("QUALITY")
636 quality.hds_type = "QUALITY"
637 quality.children["BADBITS"] = HdsPrimitive.array(np.array(255, dtype=np.uint8))
638 if "QUALITY" not in quality.children or not isinstance(quality.children["QUALITY"], HdsStructure):
639 quality.children["QUALITY"] = HdsStructure("ARRAY")
640 quality_array = quality.get_structure("QUALITY")
641 quality_array.hds_type = "ARRAY"
642 quality_array.children.setdefault("ORIGIN", HdsPrimitive.array(np.zeros(2, dtype=np.int32)))
643 quality_array.children.setdefault("BAD_PIXEL", HdsPrimitive.array(np.array(False, dtype=np.bool_)))
644 self._flush()
645 return self._file["/QUALITY/QUALITY"]
647 def _write_quality_array(self, quality: np.ndarray) -> None:
648 """Write or replace the NDF QUALITY array."""
649 self._set_quality_array(quality)
650 self._flush()
652 def _set_quality_array(self, quality: np.ndarray) -> None:
653 """Set or replace the NDF QUALITY array in the model."""
654 root = self._document.ensure_ndf("/")
655 root.set_quality(
656 NdfQuality(
657 NdfArray(
658 quality,
659 origin=np.zeros(2, dtype=np.int32),
660 bad_pixel=False,
661 compression_options=self._compression_options,
662 )
663 )
664 )
666 def _collapse_mask_to_quality(self, array: np.ndarray) -> np.ndarray:
667 """Compress an NDF-native 3-D mask array into 2-D QUALITY.
669 The input array is in HDF5 storage order ``(mask-byte, y, x)``.
670 Single-byte masks copy directly to preserve bit values. Wider masks
671 collapse to 1 where any byte is non-zero and 0 otherwise.
672 """
673 if array.shape[0] == 1:
674 return array[0, :, :]
675 return np.any(array != 0, axis=0).astype(np.uint8)
677 def _write_origin_for_array(self, struct_path: str, array: np.ndarray) -> None:
678 """Write a placeholder ORIGIN of zeros (int64).
680 The top-level `write` function overwrites this
681 with bbox-derived values via :meth:`set_array_origin` once the
682 bbox is known.
683 """
684 struct = self._document.root.get_structure(struct_path)
685 if "ORIGIN" not in struct.children:
686 struct.children["ORIGIN"] = HdsPrimitive.array(np.zeros(array.ndim, dtype=np.int64))
688 def set_array_origin(self, struct_path: str, origin: tuple[int, ...]) -> None:
689 """Overwrite the ORIGIN of an ARRAY structure.
691 Parameters
692 ----------
693 struct_path
694 HDF5 path to the ARRAY structure (e.g. ``"/DATA_ARRAY"``).
695 origin
696 Origin in NDF/Fortran axis order (e.g. ``(x_min, y_min)``
697 for a 2D image with bbox lower bound ``(x_min, y_min)``).
698 """
699 struct = self._document.root.get_structure(struct_path)
700 origin_dtype = np.int32 if struct_path == "/QUALITY/QUALITY" else np.int64
701 origin_array = np.asarray(origin, dtype=origin_dtype)
702 data_node = struct.children.get("DATA")
703 if isinstance(data_node, HdsPrimitive): 703 ↛ 707line 703 didn't jump to line 707 because the condition on line 703 was always true
704 data_ndim = data_node.read_array().ndim
705 if origin_array.size < data_ndim:
706 origin_array = np.pad(origin_array, (0, data_ndim - origin_array.size))
707 struct.children["ORIGIN"] = HdsPrimitive.array(origin_array)
708 self._flush()
710 def _ensure_model_structure(self, path: str, hds_type: str | None = None) -> HdsStructure:
711 """Return or create a structure in the NDF document model.
713 With ``hds_type=None`` the leaf type is derived from its name;
714 see `HdsStructure.ensure_structure`.
715 """
716 if path in ("", "/"): 716 ↛ 717line 716 didn't jump to line 717 because the condition on line 716 was never true
717 return self._document.root
718 return self._document.root.ensure_structure(path, hds_type)
720 def _archive_path_to_hdf5_path(self, archive_path: str) -> str:
721 """Translate an archive path to this layout's HDF5 path.
723 Parameters
724 ----------
725 archive_path
726 Serialization archive path to translate.
727 """
728 if self._lsst_path == "/MORE/LSST": 728 ↛ 730line 728 didn't jump to line 730 because the condition on line 728 was always true
729 return archive_path_to_hdf5_path(archive_path, self._name_shrinker)
730 if not archive_path:
731 return f"{self._lsst_path}/JSON"
732 components = archive_path_to_hdf5_path_components(archive_path, self._name_shrinker)
733 return f"{self._lsst_path}/{'/'.join(components)}"
735 def _register_hdf5_path(self, hdf5_path: str, logical_id: str) -> None:
736 """Record that ``logical_id`` owns ``hdf5_path``; raise on collision.
738 ``logical_id`` is the un-shrunk archive path (with any version suffix
739 applied), which is unique per logical write.
740 Two different archive entries shrinking to the same HDS path would
741 silently clobber one another, so this fails loudly instead.
743 Parameters
744 ----------
745 hdf5_path
746 HDF5 path being claimed.
747 logical_id
748 Un-shrunk archive path that owns ``hdf5_path``.
749 """
750 previous = self._hdf5_path_owners.get(hdf5_path)
751 if previous is not None and previous != logical_id:
752 raise ValueError(
753 f"NDF/HDS name collision: archive entries {previous!r} and {logical_id!r} "
754 f"both map to {hdf5_path!r} after shrinking to the {_hds.DAT__SZNAM}-character "
755 f"HDS name limit; rename one of them."
756 )
757 self._hdf5_path_owners[hdf5_path] = logical_id
759 def _versioned_archive_path(self, name: str, version: int) -> tuple[str, str]:
760 """Return ``(archive_path, logical_id)`` for a hoisted name.
762 ``archive_path`` has any version suffix applied to its leaf through the
763 version-aware shrinker (so the suffix survives the later per-component
764 shrink); ``logical_id`` is the un-shrunk version-applied path used for
765 collision detection.
766 """
767 archive_path = name if name.startswith("/") else f"/{name}"
768 if version <= 1:
769 return archive_path, archive_path
770 head, sep, leaf = archive_path.rpartition("/")
771 logical_id = f"{archive_path}_{version}"
772 versioned = f"{head}{sep}{self._name_shrinker.shrink_versioned(leaf, version)}"
773 return versioned, logical_id
775 def _has_model_path(self, path: str) -> bool:
776 """Return `True` if a path exists in the NDF document model."""
777 try:
778 self._document.get(path)
779 except KeyError:
780 return False
781 return True
783 def _model_array_ndim(self, struct_path: str) -> int:
784 """Return the dimensionality of an ARRAY structure's DATA primitive."""
785 struct = self._document.root.get_structure(struct_path)
786 data_node = struct.children["DATA"]
787 if not isinstance(data_node, HdsPrimitive): 787 ↛ 788line 787 didn't jump to line 788 because the condition on line 787 was never true
788 raise TypeError(f"{struct_path}/DATA is not an HDS primitive.")
789 return data_node.read_array().ndim
791 def _flush(self) -> None:
792 """Synchronize the Python NDF document model to the open HDF5 file."""
793 self._document.write_to_hdf5(self._file)
795 def add_table(
796 self,
797 table: astropy.table.Table,
798 *,
799 name: str | None = None,
800 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
801 ) -> TableModel:
802 columns = TableColumnModel.from_table(table, inline=True)
803 return TableModel(columns=columns, meta=table.meta)
805 def add_structured_array(
806 self,
807 array: np.ndarray,
808 *,
809 name: str | None = None,
810 units: Mapping[str, astropy.units.Unit] | None = None,
811 descriptions: Mapping[str, str] | None = None,
812 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
813 ) -> TableModel:
814 if name is None: 814 ↛ 815line 814 didn't jump to line 815 because the condition on line 814 was never true
815 columns = TableColumnModel.from_record_array(array, inline=True)
816 for c in columns:
817 if units and (unit := units.get(c.name)):
818 c.unit = unit
819 if descriptions and (description := descriptions.get(c.name)):
820 c.description = description
821 return TableModel(columns=columns)
822 name, version = self._register_name(name)
823 base_path, base_logical = self._versioned_archive_path(name, version)
824 columns = TableColumnModel.from_record_dtype(array.dtype)
825 for c in columns:
826 if len(columns) == 1:
827 archive_path = base_path
828 logical_id = base_logical
829 else:
830 archive_path = f"{base_path}/{c.name}"
831 logical_id = f"{base_logical}/{c.name}"
832 sub_ndf_path = self._archive_path_to_hdf5_path(archive_path)
833 self._register_hdf5_path(sub_ndf_path, logical_id)
834 column_array = np.asarray(array[c.name])
835 sub_ndf = self._document.ensure_ndf(sub_ndf_path)
836 sub_ndf.set_array_component(
837 "DATA_ARRAY",
838 column_array,
839 origin=np.zeros(column_array.ndim, dtype=np.int64),
840 compression_options=self._compression_options,
841 )
842 assert isinstance(c.data, ArrayReferenceModel)
843 c.data.source = f"ndf:{sub_ndf_path}/DATA_ARRAY/DATA"
844 for c in columns:
845 if units and (unit := units.get(c.name)):
846 c.unit = unit
847 if descriptions and (description := descriptions.get(c.name)):
848 c.description = description
849 self._flush()
850 return TableModel(columns=columns)