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