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

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. 

11 

12from __future__ import annotations 

13 

14__all__ = ( 

15 "NdfOutputArchive", 

16 "write", 

17) 

18 

19import os 

20from collections.abc import Callable, Hashable, Iterator, Mapping, Sequence 

21from contextlib import contextmanager 

22from pathlib import Path 

23from typing import Any, Self 

24 

25import astropy.io.fits 

26import astropy.table 

27import astropy.units 

28import h5py 

29import numpy as np 

30import pydantic 

31 

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 

59 

60_NDF_FORMAT_VERSION = 1 

61"""Container layout version for files written by `NdfOutputArchive`. 

62 

63Bumps when the NDF layout (NdfDocument shape, .MORE.LSST contents) 

64changes. Independent of any data-model ``SCHEMA_VERSION``. 

65""" 

66 

67 

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. 

77 

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. 

98 

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 

124 

125 

126def _origin_from_bbox(bbox: Any) -> tuple[int, ...]: 

127 """Extract NDF/Fortran-order origin tuple from an lsst.images Box. 

128 

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 ) 

145 

146 

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() 

153 

154 

155def _fits_header_records(header: astropy.io.fits.Header) -> list[str]: 

156 """Return fixed-width FITS records for an opaque NDF FITS extension. 

157 

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)] 

171 

172 

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 {} 

187 

188 

189def _show_ast_for_ndf(ast_frame_set: Any, bbox: Any | None) -> str: 

190 """Return AST Channel text matching Starlink NDF WCS serialization. 

191 

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) 

207 

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") 

213 

214 stream = StringStream() 

215 channel = Channel(stream, options="Full=-1,Comment=0,Indent=0") 

216 channel.write(ast_frame_set) 

217 return stream.getSinkData() 

218 

219 

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. 

227 

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) 

238 

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 ) 

251 

252 stream = StringStream() 

253 channel = Channel(stream, options="Full=-1,Comment=0,Indent=0") 

254 channel.write(ast_frame_set) 

255 return stream.getSinkData() 

256 

257 

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. 

262 

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 """ 

286 

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() 

313 

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`. 

328 

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). 

332 

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() 

373 

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. 

384 

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. 

393 

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() 

432 

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))) 

451 

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) 

457 

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 

483 

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 

490 

491 def iter_frame_sets(self) -> Iterator[tuple[FrameSet, NdfPointerModel]]: 

492 return iter(self._frame_sets) 

493 

494 _COMPATIBLE_MASK_DTYPES = (np.dtype(np.uint8),) 

495 _prefer_native_mask_arrays = True 

496 

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 ) 

593 

594 def _ensure_path(self, path: str) -> h5py.Group: 

595 """Walk/create groups for an HDF5 absolute path. 

596 

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] 

604 

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] 

610 

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") 

614 

615 def _ensure_quality_structure(self) -> h5py.Group: 

616 """Ensure ``/QUALITY`` exists with ``CLASS="QUALITY"`` and BADBITS. 

617 

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"] 

629 

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"] 

646 

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() 

651 

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 ) 

665 

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. 

668 

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) 

676 

677 def _write_origin_for_array(self, struct_path: str, array: np.ndarray) -> None: 

678 """Write a placeholder ORIGIN of zeros (int64). 

679 

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)) 

687 

688 def set_array_origin(self, struct_path: str, origin: tuple[int, ...]) -> None: 

689 """Overwrite the ORIGIN of an ARRAY structure. 

690 

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() 

709 

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. 

712 

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) 

719 

720 def _archive_path_to_hdf5_path(self, archive_path: str) -> str: 

721 """Translate an archive path to this layout's HDF5 path. 

722 

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)}" 

734 

735 def _register_hdf5_path(self, hdf5_path: str, logical_id: str) -> None: 

736 """Record that ``logical_id`` owns ``hdf5_path``; raise on collision. 

737 

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. 

742 

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 

758 

759 def _versioned_archive_path(self, name: str, version: int) -> tuple[str, str]: 

760 """Return ``(archive_path, logical_id)`` for a hoisted name. 

761 

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 

774 

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 

782 

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 

790 

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) 

794 

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) 

804 

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)