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