Coverage for python/lsst/images/ndf/_input_archive.py: 15%
294 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-06 08:37 +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__ = ("NdfInputArchive", "read_starlink")
16import logging
17from collections.abc import Callable, Iterator
18from contextlib import contextmanager
19from types import EllipsisType
20from typing import Any, Self
22import astropy.io.fits
23import astropy.table
24import astropy.units as u
25import h5py
26import numpy as np
28from lsst.resources import ResourcePath, ResourcePathExpression
30from .._geom import Box
31from .._image import Image
32from .._mask import Mask, MaskPlane, MaskSchema
33from .._masked_image import MaskedImage
34from .._transforms import FrameSet, Projection
35from .._transforms import _ast as astshim
36from .._transforms._frames import GeneralFrame
37from ..fits._common import FitsOpaqueMetadata
38from ..serialization import (
39 ArchiveInfo,
40 ArchiveReadError,
41 ArchiveTree,
42 ArrayReferenceModel,
43 InlineArrayModel,
44 InputArchive,
45 TableModel,
46 no_header_updates,
47 parameterize_tree,
48)
49from ..serialization._common import _check_format_version
50from . import _hds
51from ._common import NdfPointerModel
52from ._model import HdsPrimitive, NdfDocument
54_LOG = logging.getLogger(__name__)
56_NDF_FORMAT_VERSION = 1
57"""Container layout version this release of `NdfInputArchive` understands."""
60class NdfInputArchive(InputArchive[NdfPointerModel]):
61 """Reads HDS-on-HDF5 NDF files written by `NdfOutputArchive`.
63 Instances should only be constructed via the :meth:`open` context
64 manager.
66 Parameters
67 ----------
68 file
69 Open `h5py.File` handle. Owned by the caller of :meth:`open`;
70 the archive does not close it.
71 """
73 def __init__(self, file: h5py.File) -> None:
74 self._file = file
75 self._document = NdfDocument.from_hdf5(file)
76 self._opaque_metadata = FitsOpaqueMetadata()
77 self._deserialized_pointer_cache: dict[str, Any] = {}
78 self._frame_set_cache: dict[str, FrameSet] = {}
79 self._read_opaque_fits_metadata()
80 self._check_format_version()
82 @classmethod
83 def get_basic_info(cls, path: ResourcePathExpression) -> ArchiveInfo:
84 """Read the schema URL from the ``DATA_MODEL`` scalar and the
85 ``FORMAT_VERSION`` primitive without deserializing pixel data.
87 Both live at the fixed location ``/MORE/LSST`` (or ``/LSST``); we read
88 only those and never search at arbitrary depth, so nested pointer
89 trees cannot be mistaken for the top level. Reading ``DATA_MODEL``
90 directly avoids parsing the (potentially large) JSON tree.
91 """
92 ospath = ResourcePath(path).ospath
93 schema_url: str | None = None
94 format_version = 1
96 with h5py.File(ospath, "r") as handle:
97 for prefix in ("MORE/LSST", "LSST"):
98 data_model = handle.get(f"{prefix}/DATA_MODEL")
99 if not isinstance(data_model, h5py.Dataset):
100 continue
101 schema_url = np.asarray(data_model).tobytes().decode("ascii").rstrip("\x00").strip()
102 fmt_node = handle.get(f"{prefix}/FORMAT_VERSION")
103 if fmt_node is not None:
104 format_version = int(np.asarray(fmt_node).item())
105 break
106 if not schema_url:
107 raise ArchiveReadError(
108 f"Could not read the schema of {path!r} from /MORE/LSST/DATA_MODEL or /LSST/DATA_MODEL."
109 )
110 return ArchiveInfo.from_schema_url(schema_url, format_version=format_version)
112 @classmethod
113 @contextmanager
114 def open_tree(
115 cls,
116 path: ResourcePathExpression,
117 tree_cls: type[ArchiveTree],
118 *,
119 partial: bool = True,
120 **backend_kwargs: Any,
121 ) -> Iterator[tuple[Self, ArchiveTree]]:
122 """Open the NDF file and yield ``(archive, tree)``.
124 Requires the symmetric LSST JSON tree; ``partial`` is accepted but
125 not meaningful, since h5py reads lazily regardless.
126 """
127 parameterized = parameterize_tree(tree_cls, NdfPointerModel)
128 with cls.open(path) as archive:
129 if archive._get_main_json_path() is None:
130 raise ArchiveReadError(
131 f"{path!r} has no LSST JSON tree; only the symmetric read path is supported."
132 )
133 tree = archive.get_tree(parameterized)
134 yield archive, tree
136 @classmethod
137 @contextmanager
138 def open(cls, path: ResourcePathExpression) -> Iterator[Self]:
139 """Open an NDF file for reading and yield an `NdfInputArchive`.
141 Remote ResourcePaths are materialised locally first; fsspec-direct
142 h5py reads are a deferred follow-up.
143 """
144 rp = ResourcePath(path)
145 with rp.as_local() as local:
146 with h5py.File(local.ospath, "r") as f:
147 yield cls(f)
149 def get_tree[T: ArchiveTree](self, model_type: type[T]) -> T:
150 """Read and validate the main Pydantic tree at ``/MORE/LSST/JSON``."""
151 json_path = self._get_main_json_path()
152 if json_path is None:
153 raise ArchiveReadError(
154 "File has no /MORE/LSST/JSON tree; this is either a "
155 "Starlink-only NDF (use ndf.read_starlink() for auto-detect) or "
156 "the file was written by an unrelated tool."
157 )
158 json_text = _read_json_record(self._get_primitive(json_path), json_path)
159 return model_type.model_validate_json(json_text)
161 def deserialize_pointer[U: ArchiveTree, V](
162 self,
163 pointer: NdfPointerModel,
164 model_type: type[U],
165 deserializer: Callable[[U, InputArchive[NdfPointerModel]], V],
166 ) -> V:
167 # Cache by pointer.path so repeated dereferences reuse the same
168 # deserialised result and don't re-run the deserializer.
169 if (cached := self._deserialized_pointer_cache.get(pointer.path)) is not None:
170 return cached
171 if not self._has_model_path(pointer.path):
172 raise ArchiveReadError(f"Pointer reference {pointer.path!r} not found in NDF file.")
173 primitive = self._get_primitive(pointer.path)
174 json_text = _read_json_record(primitive, pointer.path)
175 model = model_type.model_validate_json(json_text)
176 result = deserializer(model, self)
177 self._deserialized_pointer_cache[pointer.path] = result
178 if isinstance(result, FrameSet):
179 self._frame_set_cache[pointer.path] = result
180 return result
182 def get_frame_set(self, pointer: NdfPointerModel) -> FrameSet:
183 try:
184 return self._frame_set_cache[pointer.path]
185 except KeyError:
186 raise AssertionError(
187 f"Frame set at {pointer.path!r} must be deserialised via "
188 f"deserialize_pointer before any dependent transform can be."
189 ) from None
191 def get_array(
192 self,
193 model: ArrayReferenceModel | InlineArrayModel,
194 *,
195 slices: tuple[slice, ...] | EllipsisType = ...,
196 strip_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
197 ) -> np.ndarray:
198 if isinstance(model, InlineArrayModel):
199 data: np.ndarray = np.array(model.data, dtype=model.datatype.to_numpy())
200 return data if slices is ... else data[slices]
201 if not isinstance(model.source, str) or not model.source.startswith("ndf:"):
202 raise ArchiveReadError(
203 f"NdfInputArchive cannot resolve array source {model.source!r}; "
204 f"expected an 'ndf:<HDF5-path>' reference."
205 )
206 path = model.source[len("ndf:") :]
207 if not self._has_model_path(path):
208 raise ArchiveReadError(f"Array reference {path!r} not in file.")
209 primitive = self._get_primitive(path)
210 # h5py supports lazy slicing via dataset[slices].
211 if isinstance(primitive.data, h5py.Dataset):
212 return primitive.data[()] if slices is ... else primitive.data[slices]
213 data = primitive.read_array()
214 return data if slices is ... else data[slices]
216 def get_table(
217 self,
218 model: TableModel,
219 strip_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
220 ) -> astropy.table.Table:
221 result = astropy.table.Table(meta=model.meta)
222 for column_model in model.columns:
223 if isinstance(column_model.data, InlineArrayModel):
224 data: Any = column_model.data.data
225 else:
226 data = self.get_array(column_model.data, strip_header=strip_header)
227 result[column_model.name] = astropy.table.Column(
228 data,
229 name=column_model.name,
230 dtype=column_model.data.datatype.to_numpy(),
231 unit=column_model.unit,
232 description=column_model.description,
233 meta=column_model.meta,
234 )
235 return result
237 def get_structured_array(
238 self,
239 model: TableModel,
240 strip_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
241 ) -> np.ndarray:
242 return self.get_table(model, strip_header).as_array()
244 def _read_opaque_fits_metadata(self) -> None:
245 if not self._has_model_path("/MORE/FITS"):
246 return
247 cards = self._get_primitive("/MORE/FITS").read_char_array()
248 # FITS Header.fromstring expects fixed-width 80-char cards
249 # concatenated; pad each card defensively so readers tolerate
250 # files written with shorter widths.
251 header = astropy.io.fits.Header.fromstring("".join(c.ljust(80) for c in cards))
252 self._opaque_metadata.add_header(header, name="", ver=1)
254 def get_opaque_metadata(self) -> FitsOpaqueMetadata:
255 return self._opaque_metadata
257 def _get_main_json_path(self) -> str | None:
258 """Return the path of the main LSST JSON tree, if present."""
259 for path in ("/MORE/LSST/JSON", "/LSST/JSON"):
260 if self._has_model_path(path):
261 return path
262 return None
264 def _check_format_version(self) -> None:
265 """Read FORMAT_VERSION from the NDF top-level structure and check it.
267 Absence is treated as ``1`` (legacy default). DATA_MODEL is
268 informational only on read; the JSON tree's ``schema_version`` /
269 ``min_read_version`` drive data-model compatibility.
270 """
271 on_disk = 1
272 for prefix in ("/MORE/LSST", "/LSST"):
273 path = f"{prefix}/FORMAT_VERSION"
274 if self._has_model_path(path):
275 primitive = self._get_primitive(path)
276 # We wrote the version as a 0-d int32 numpy array; .item()
277 # unwraps to a Python int.
278 on_disk = int(primitive.read_array().item())
279 break
280 _check_format_version("ndf", on_disk, _NDF_FORMAT_VERSION)
282 def _has_model_path(self, path: str) -> bool:
283 """Return `True` if a path exists in the NDF document model."""
284 try:
285 self._document.get(path)
286 except KeyError:
287 return False
288 return True
290 def _get_primitive(self, path: str) -> HdsPrimitive:
291 """Return a primitive component from the NDF document model."""
292 node = self._document.get(path)
293 if not isinstance(node, HdsPrimitive):
294 raise ArchiveReadError(f"NDF reference {path!r} is not a primitive dataset.")
295 return node
298def read_starlink[T: Any](cls: type[T], path: ResourcePathExpression) -> T:
299 """Reconstruct an `~lsst.images.Image` or `~lsst.images.MaskedImage`
300 from a schema-less Starlink NDF.
302 Files written by this package carry a ``/MORE/LSST/JSON`` tree and are
303 read through the generic `lsst.images.serialization.read` /
304 `lsst.images.serialization.open`. A Starlink-produced NDF has no such
305 tree and therefore no schema, so it cannot go through that path; this
306 function auto-detects a minimal recognised-component set
307 (``DATA_ARRAY``, ``VARIANCE``, ``QUALITY``, ``MORE.FITS``) instead.
308 ``WCS`` is reconstructed when possible; other components are
309 logged-and-dropped.
311 Parameters
312 ----------
313 cls
314 Expected return type; `~lsst.images.Image` and
315 `~lsst.images.MaskedImage` are the only types the auto-detect path
316 can produce.
317 path
318 File path or `lsst.resources.ResourcePathExpression`.
320 Returns
321 -------
322 object
323 The deserialized ``cls`` instance.
325 Raises
326 ------
327 ArchiveReadError
328 If the file has an LSST JSON tree (use the generic ``read`` instead)
329 or no recognised ``DATA_ARRAY`` component.
330 """
331 with NdfInputArchive.open(path) as archive:
332 if archive._get_main_json_path() is not None:
333 raise ArchiveReadError(
334 f"{path!r} has an LSST JSON tree; read it with serialization.read()/open()."
335 )
336 return _read_auto_detect(cls, archive)
339def _read_auto_detect[T: Any](cls: type[T], archive: NdfInputArchive) -> T:
340 """Reconstruct an `Image` (or `MaskedImage`) from a Starlink NDF.
342 Recognised components: ``DATA_ARRAY`` (in either simple or complex
343 form), ``VARIANCE``, ``QUALITY``, ``MORE.FITS``. Other components
344 (``WCS``, ``HISTORY``, ``AXIS``, ``LABEL``, custom ``MORE.*``,
345 ``_LOGICAL`` primitives) are warned-and-dropped.
346 """
347 f = archive._file
348 ndf_group = _locate_ndf_root(f)
350 # DATA_ARRAY is required.
351 if "DATA_ARRAY" not in ndf_group:
352 raise ArchiveReadError(f"Auto-detect read of {f.filename!r}: no DATA_ARRAY component.")
353 data_arr, bbox = _read_data_array_with_bbox(ndf_group["DATA_ARRAY"])
355 # VARIANCE / QUALITY are optional.
356 variance_arr: np.ndarray | None = None
357 variance_bbox: Any | None = None
358 if "VARIANCE" in ndf_group:
359 variance_arr, variance_bbox = _read_data_array_with_bbox(ndf_group["VARIANCE"])
360 quality_arr: np.ndarray | None = None
361 quality_bbox: Any | None = None
362 quality_badbits = 255
363 if "QUALITY" in ndf_group and isinstance(ndf_group["QUALITY"], h5py.Group):
364 q = ndf_group["QUALITY"]
365 quality_badbits = _read_quality_badbits(q)
366 if "QUALITY" in q and isinstance(q["QUALITY"], h5py.Dataset):
367 quality_arr = _validate_quality_array(_hds.read_array(q["QUALITY"]))
368 quality_bbox = _make_bbox(x_min=0, y_min=0, array=quality_arr)
369 elif "QUALITY" in q and isinstance(q["QUALITY"], h5py.Group):
370 quality_arr, quality_bbox = _read_data_array_with_bbox(q["QUALITY"])
371 quality_arr = _validate_quality_array(quality_arr)
373 projection: Projection | None = None
374 if "WCS" in ndf_group:
375 try:
376 wcs_group = ndf_group["WCS"]
377 if isinstance(wcs_group, h5py.Group) and "DATA" in wcs_group:
378 wcs_lines = _hds.read_char_array(wcs_group["DATA"])
379 wcs_text = _hds.decode_ndf_ast_data(wcs_lines)
380 ast_obj = astshim.Object.fromString(wcs_text)
381 if isinstance(ast_obj, astshim.FrameSet):
382 pixel_frame = GeneralFrame(unit=u.pix)
383 projection = Projection.from_ast_frame_set(
384 ast_obj,
385 pixel_frame,
386 pixel_bounds=bbox,
387 )
388 except Exception:
389 _LOG.warning(
390 "Could not reconstruct Projection from WCS in %s; dropping.",
391 f.filename,
392 exc_info=True,
393 )
395 unit = _read_ndf_units(ndf_group)
397 # Anything unrecognised: warn-and-drop.
398 recognised = {
399 "DATA_ARRAY",
400 "VARIANCE",
401 "QUALITY",
402 "WCS",
403 "MORE",
404 "TITLE",
405 "LABEL",
406 "UNITS",
407 "HISTORY",
408 "AXIS",
409 }
410 for name in ndf_group:
411 if name not in recognised:
412 _LOG.warning(
413 "Ignoring unrecognised NDF component %s/%s during auto-detect read.",
414 ndf_group.name,
415 name,
416 )
418 # Build the requested in-memory object. Any NDF can be read as an Image;
419 # MaskedImage construction uses whatever VARIANCE/QUALITY are present and
420 # lets the MaskedImage constructor provide defaults for missing planes.
421 image = Image(data_arr, bbox=bbox, unit=unit, projection=projection)
422 obj: Any
423 if cls is Image:
424 obj = image
425 elif issubclass(cls, MaskedImage):
426 if quality_arr is not None:
427 schema = _make_quality_mask_schema(quality_badbits)
428 mask = Mask(quality_arr[:, :, np.newaxis], schema=schema, bbox=quality_bbox)
429 else:
430 schema = MaskSchema([MaskPlane(name="BAD", description="Bad pixel.")])
431 mask = None
432 variance = Image(variance_arr, bbox=variance_bbox) if variance_arr is not None else None
433 obj = cls(
434 image=image,
435 mask=mask,
436 mask_schema=schema if mask is None else None,
437 variance=variance,
438 )
439 else:
440 raise ArchiveReadError(
441 f"Auto-detect can produce Image or MaskedImage, but caller asked for {cls.__name__}."
442 )
443 obj._opaque_metadata = archive.get_opaque_metadata()
444 return obj
447def _read_ndf_units(ndf_group: h5py.Group) -> u.UnitBase | None:
448 """Read the NDF UNITS component, if present."""
449 if "UNITS" not in ndf_group or not isinstance(ndf_group["UNITS"], h5py.Dataset):
450 return None
451 dataset = ndf_group["UNITS"]
452 if dataset.dtype.kind != "S":
453 _LOG.warning("Ignoring non-character NDF UNITS component in %s.", ndf_group.name)
454 return None
455 if dataset.ndim == 0:
456 raw = dataset[()]
457 if isinstance(raw, np.bytes_):
458 raw = bytes(raw)
459 if not isinstance(raw, bytes):
460 return None
461 units_text = raw.decode("ascii").rstrip(" ")
462 else:
463 records = _hds.read_char_array(dataset)
464 units_text = records[0] if records else ""
465 if not units_text:
466 return None
467 for kwargs in ({"format": "fits"}, {}):
468 try:
469 return u.Unit(units_text, **kwargs)
470 except ValueError:
471 continue
472 _LOG.warning("Could not parse NDF UNITS value %r in %s.", units_text, ndf_group.name)
473 return None
476def _read_quality_badbits(quality_group: h5py.Group) -> int:
477 """Read the scalar NDF QUALITY.BADBITS value."""
478 badbits = quality_group.get("BADBITS")
479 if not isinstance(badbits, h5py.Dataset):
480 return 255
481 value = np.asarray(_hds.read_array(badbits)).reshape(-1)
482 if value.size == 0:
483 return 255
484 return int(value[0])
487def _validate_quality_array(quality: np.ndarray) -> np.ndarray:
488 """Return an NDF QUALITY array as a `numpy.uint8` mask plane."""
489 if quality.dtype != np.dtype(np.uint8):
490 raise ArchiveReadError(f"NDF QUALITY array has dtype {quality.dtype}; expected uint8.")
491 return quality
494def _make_quality_mask_schema(badbits: int) -> MaskSchema:
495 """Create a fallback `MaskSchema` for an unnamed 8-bit QUALITY array."""
496 planes = []
497 for bit in range(8):
498 mask = 1 << bit
499 description = f"NDF QUALITY bit {bit}."
500 if badbits & mask:
501 description += " Selected by BADBITS."
502 planes.append(MaskPlane(name=f"MASK{bit}", description=description))
503 return MaskSchema(planes, dtype=np.uint8)
506def _locate_ndf_root(f: h5py.File) -> h5py.Group:
507 """Return the group representing the top-level NDF.
509 Most files have the NDF at the root group itself. A few wrap it
510 in a single-child container at the root; we accept that shape
511 too. Anything more elaborate raises.
512 """
513 root_class = f["/"].attrs.get(_hds.ATTR_CLASS)
514 if isinstance(root_class, bytes):
515 root_class = root_class.decode("ascii")
516 if root_class == "NDF":
517 return f["/"]
518 # Maybe a one-level container.
519 candidates = []
520 for name, child in f["/"].items():
521 if isinstance(child, h5py.Group):
522 cls_attr = child.attrs.get(_hds.ATTR_CLASS)
523 if isinstance(cls_attr, bytes):
524 cls_attr = cls_attr.decode("ascii")
525 if cls_attr == "NDF":
526 candidates.append(name)
527 if len(candidates) == 1:
528 return f[candidates[0]]
529 raise ArchiveReadError(
530 f"Could not locate top-level NDF in {f.filename!r}; "
531 f"expected the root group or a single NDF-typed child."
532 )
535def _read_data_array_with_bbox(
536 obj: h5py.Group | h5py.Dataset,
537) -> tuple[np.ndarray, Any]:
538 """Read a DATA_ARRAY component in either simple or complex form.
540 The complex form (what our writer always produces) is an HDS
541 ARRAY structure (h5py group with CLASS="ARRAY") containing
542 ``DATA`` and ``ORIGIN`` primitives. The simple form is a bare
543 primitive dataset.
545 Returns
546 -------
547 array, bbox : tuple
548 ``array`` is the C-order numpy data (shape ``(height, width)``
549 for 2D images). ``bbox`` is constructed from the ORIGIN if
550 present, else from a default origin of (0, 0).
551 """
552 if isinstance(obj, h5py.Dataset):
553 # Simple form.
554 array = _hds.read_array(obj)
555 bbox = _make_bbox(x_min=0, y_min=0, array=array)
556 return array, bbox
557 # Complex form: an HDS structure with DATA + ORIGIN.
558 data = _hds.read_array(obj["DATA"])
559 if "ORIGIN" in obj:
560 origin = _hds.read_array(obj["ORIGIN"])
561 bbox = _make_bbox(x_min=int(origin[0]), y_min=int(origin[1]), array=data)
562 else:
563 bbox = _make_bbox(x_min=0, y_min=0, array=data)
564 return data, bbox
567def _read_json_record(primitive: HdsPrimitive, path: str) -> str:
568 """Read a JSON document stored as a single _CHAR*N record.
570 Our writer always emits JSON trees as a single-element character
571 array sized to the document. Joining multiple records would lose
572 trailing whitespace inside JSON string values, since
573 `read_char_array` strips trailing spaces per record.
574 """
575 records = primitive.read_char_array()
576 if len(records) != 1:
577 raise ArchiveReadError(f"Expected a single _CHAR*N record at {path!r}, got {len(records)}.")
578 return records[0]
581def _make_bbox(*, x_min: int, y_min: int, array: np.ndarray) -> Any:
582 """Build an lsst.images.Box for a 2D image array.
584 The array is C-order ``(height, width)``. NDF stores ``ORIGIN``
585 in Fortran axis order ``(x_min, y_min)``.
586 """
587 if array.ndim != 2:
588 raise ArchiveReadError(f"Auto-detect read only supports 2D arrays, got ndim={array.ndim}.")
589 # Box.from_shape takes (height, width) and start=(y_start, x_start).
590 return Box.from_shape(array.shape, start=(y_min, x_min))