Coverage for python/lsst/images/ndf/_output_archive.py: 81%

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 typing import Any, Self 

23 

24import astropy.io.fits 

25import astropy.table 

26import astropy.units 

27import h5py 

28import numpy as np 

29import pydantic 

30 

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 ( 

52 HdsNameShrinker, 

53 NdfPointerModel, 

54 archive_path_to_hdf5_path, 

55 archive_path_to_hdf5_path_components, 

56) 

57from ._model import HdsPrimitive, HdsStructure, Ndf, NdfArray, NdfContainer, NdfDocument, NdfQuality, NdfWcs 

58 

59_NDF_FORMAT_VERSION = 1 

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

61 

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

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

64""" 

65 

66 

67def write( 

68 obj: Any, 

69 filename: str | None = None, 

70 *, 

71 metadata: dict[str, MetadataValue] | None = None, 

72 butler_info: ButlerInfo | None = None, 

73 compression_options: Mapping[str, Any] | None = None, 

74) -> ArchiveTree: 

75 """Write a serializable object to an NDF (HDS-on-HDF5) file. 

76 

77 Parameters 

78 ---------- 

79 obj 

80 Object with a ``serialize`` method. May carry an 

81 ``_opaque_metadata`` attribute (a 

82 `~lsst.images.fits.FitsOpaqueMetadata`) 

83 whose primary-HDU header gets written to ``/MORE/FITS``. This 

84 preserves FITS cards on objects that originated from a FITS read; 

85 butler provenance is conveyed through ``butler_info`` instead. 

86 filename 

87 Path to write to. Must not already exist. If `None`, an 

88 in-memory HDF5 file is used and the on-disk artefact is 

89 discarded; the returned tree still reflects all the writes the 

90 archive made (useful for tests). 

91 metadata, butler_info 

92 Optional caller-supplied entries that are written into the 

93 returned `~lsst.images.serialization.ArchiveTree`. 

94 compression_options 

95 Optional dict forwarded to the archive constructor for h5py 

96 dataset compression. 

97 

98 Returns 

99 ------- 

100 `~lsst.images.serialization.ArchiveTree` 

101 The Pydantic tree the object's ``serialize`` produced (with 

102 ``metadata``/``butler_info`` applied). 

103 """ 

104 opaque_metadata = getattr(obj, "_opaque_metadata", None) 

105 if not isinstance(opaque_metadata, FitsOpaqueMetadata): 

106 opaque_metadata = FitsOpaqueMetadata() 

107 archive_default_name = getattr(obj, "_archive_default_name", None) 

108 with NdfOutputArchive.open( 

109 filename, 

110 compression_options=compression_options, 

111 opaque_metadata=opaque_metadata, 

112 **_get_archive_layout(obj), 

113 ) as archive: 

114 if archive_default_name is not None: 

115 tree = archive.serialize_direct(archive_default_name, obj.serialize) 

116 else: 

117 tree = obj.serialize(archive) 

118 if metadata is not None: 

119 tree.metadata.update(metadata) 

120 if butler_info is not None: 

121 tree.butler_info = butler_info 

122 archive.add_tree( 

123 tree, 

124 sky_projection=getattr(obj, "sky_projection", None), 

125 bbox=getattr(obj, "bbox", None), 

126 unit=getattr(obj, "unit", None), 

127 root_name=archive_default_name or type(obj).__name__, 

128 ) 

129 return tree 

130 

131 

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

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

134 

135 The exact attribute names on Box depend on the version. Inspect the 

136 object and pick whatever exposes the integer lower bounds. For a 2D 

137 image with bbox lower bound (x_min, y_min) the returned tuple is 

138 ``(x_min, y_min)``. 

139 """ 

140 # Box exposes .x and .y properties returning Interval objects with a 

141 # .start attribute (the lower bound). 

142 if hasattr(bbox, "x") and hasattr(bbox, "y"): 142 ↛ 147line 142 didn't jump to line 147 because the condition on line 142 was always true

143 x = bbox.x 

144 y = bbox.y 

145 if hasattr(x, "start") and hasattr(y, "start"): 145 ↛ 147line 145 didn't jump to line 147 because the condition on line 145 was always true

146 return (int(x.start), int(y.start)) 

147 raise AttributeError( 

148 f"Don't know how to extract origin from bbox of type {type(bbox).__name__!r}; " 

149 "_origin_from_bbox needs updating." 

150 ) 

151 

152 

153def _unit_to_ndf_string(unit: astropy.units.UnitBase) -> str: 

154 """Return an ASCII unit string for the NDF UNITS component.""" 

155 try: 

156 return unit.to_string(format="fits") 

157 except ValueError: 

158 return unit.to_string() 

159 

160 

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

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

163 

164 NDF ``.MORE.FITS`` is a ``_CHAR*80`` vector. Use Astropy's FITS 

165 header serializer to preserve multi-record logical cards such as 

166 CONTINUE strings, but omit the FITS END card and 2880-byte padding 

167 because this is an NDF extension component, not a complete FITS 

168 header block. 

169 """ 

170 block = header.tostring(sep="", endcard=False, padding=False) 

171 encoded = block.encode("ascii") 

172 if len(encoded) % 80: 172 ↛ 173line 172 didn't jump to line 173 because the condition on line 172 was never true

173 raise ValueError( 

174 f"FITS header block is {len(encoded)} bytes, not a multiple of the 80-byte FITS record size." 

175 ) 

176 return [block[n : n + 80] for n in range(0, len(block), 80)] 

177 

178 

179def _get_archive_layout(obj: Any) -> dict[str, Any]: 

180 """Return NDF document layout options for a top-level object.""" 

181 if isinstance(obj, ColorImage): 

182 return { 

183 "root": NdfContainer(), 

184 "lsst_path": "/LSST", 

185 "direct_ndf_array_paths": { 

186 "red": "/RED", 

187 "green": "/GREEN", 

188 "blue": "/BLUE", 

189 }, 

190 "wcs_ndf_paths": ("/RED", "/GREEN", "/BLUE"), 

191 } 

192 return {} 

193 

194 

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

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

197 

198 Tags the original base frame with ``Domain="PIXEL"`` and prepends a 

199 new ``Domain="GRID"`` base frame related to it by a `ShiftMap` whose 

200 shift converts ``bbox``-origin pixel coordinates into 1-based grid 

201 coordinates. The result is written via an abstraction-layer 

202 ``Channel`` configured with the same options the Starlink C writer 

203 uses (``Full=-1,Comment=0``; see ``ndf1Wwrt.c``) plus ``Indent=0`` so 

204 each line is just the bare AST item with the single-space prefix 

205 that ``ndf1Rdast`` strips back off on read. 

206 """ 

207 if bbox is None: 207 ↛ 208line 207 didn't jump to line 208 because the condition on line 207 was never true

208 x_shift = 1.0 

209 y_shift = 1.0 

210 else: 

211 x_shift = 1.0 - float(bbox.x.start) 

212 y_shift = 1.0 - float(bbox.y.start) 

213 

214 saved_current = ast_frame_set.current 

215 ast_frame_set.current = ast_frame_set.base 

216 ast_frame_set.domain = "PIXEL" 

217 ast_frame_set.current = saved_current 

218 _prepend_ast_shift(ast_frame_set, x=x_shift, y=y_shift, ast_domain="GRID") 

219 

220 stream = StringStream() 

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

222 channel.write(ast_frame_set) 

223 return stream.getSinkData() 

224 

225 

226def _show_mask_ast_for_ndf( 

227 parent_ast_frame_set: Any, 

228 origin: Sequence[int], 

229 *, 

230 labels: Sequence[str] = (), 

231) -> str: 

232 """Return an NDF WCS for the 3D native mask sub-NDF. 

233 

234 The first two axes reuse the parent image's pixel-to-sky mapping. The 

235 third axis is a generic mask-byte coordinate that passes through unchanged. 

236 """ 

237 n_axes = len(origin) 

238 ast_frame_set = AstFrameSet(AstFrame(n_axes, "Domain=GRID")) 

239 pixel_frame = AstFrame(n_axes, "Domain=PIXEL") 

240 for axis, label in enumerate(labels[:n_axes], start=1): 

241 pixel_frame.setLabel(axis, label) 

242 shifts = [1.0 - float(axis_origin) for axis_origin in origin] 

243 ast_frame_set.addFrame(AstFrameSet.BASE, ShiftMap(shifts), pixel_frame) 

244 

245 parent_pixel_to_sky = parent_ast_frame_set.getMapping( 

246 parent_ast_frame_set.base, 

247 parent_ast_frame_set.current, 

248 ) 

249 parent_sky_frame = parent_ast_frame_set.getFrame(parent_ast_frame_set.current) 

250 mask_axis_frame = AstFrame(1, "Domain=MASK") 

251 mask_axis_frame.setLabel(1, labels[2] if len(labels) > 2 else "mask-byte") 

252 ast_frame_set.addFrame( 

253 ast_frame_set.current, 

254 CmpMap(parent_pixel_to_sky, UnitMap(1), False), 

255 CmpFrame(parent_sky_frame, mask_axis_frame), 

256 ) 

257 

258 stream = StringStream() 

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

260 channel.write(ast_frame_set) 

261 return stream.getSinkData() 

262 

263 

264class NdfOutputArchive(OutputArchive[NdfPointerModel]): 

265 """An `~lsst.images.serialization.OutputArchive` implementation 

266 that writes HDS-on-HDF5 files compatible with the Starlink NDF data 

267 model. 

268 

269 Parameters 

270 ---------- 

271 file 

272 An open `h5py.File` opened in a writable mode. The archive does 

273 not close the file; the caller is responsible for that. 

274 compression_options 

275 Optional dict passed through to `h5py.Group.create_dataset` for image 

276 arrays (e.g. ``{"compression": "gzip", "compression_opts": 4}``). 

277 opaque_metadata 

278 Optional `~lsst.images.fits.FitsOpaqueMetadata`; if its primary-HDU 

279 header is non-empty its cards will be written to ``/MORE/FITS`` by the 

280 top-level `write` function. 

281 """ 

282 

283 def __init__( 

284 self, 

285 file: h5py.File, 

286 compression_options: Mapping[str, Any] | None = None, 

287 opaque_metadata: FitsOpaqueMetadata | None = None, 

288 root: Ndf | NdfContainer | None = None, 

289 lsst_path: str = "/MORE/LSST", 

290 direct_ndf_array_paths: Mapping[str, str] | None = None, 

291 wcs_ndf_paths: Sequence[str] = ("/",), 

292 ) -> None: 

293 super().__init__() 

294 self._file = file 

295 self._document = NdfDocument(root=root if root is not None else Ndf()) 

296 self._lsst_path = lsst_path.rstrip("/") or "/LSST" 

297 self._direct_ndf_array_paths = dict(direct_ndf_array_paths) if direct_ndf_array_paths else {} 

298 self._wcs_ndf_paths = tuple(wcs_ndf_paths) 

299 self._bbox_array_struct_paths: set[str] = set() 

300 self._compression_options = dict(compression_options) if compression_options else {} 

301 self._opaque_metadata = opaque_metadata if opaque_metadata is not None else FitsOpaqueMetadata() 

302 self._frame_sets: list[tuple[FrameSet, NdfPointerModel]] = [] 

303 self._pointers: dict[Hashable, NdfPointerModel] = {} 

304 self._hdf5_path_owners: dict[str, str] = {} 

305 self._name_shrinker = HdsNameShrinker() 

306 # Keep the open file in sync so existing direct-archive tests can 

307 # inspect it immediately, while all mutations go through the IR. 

308 self._flush() 

309 

310 @classmethod 

311 @contextmanager 

312 def open( 

313 cls, 

314 filename: str | None, 

315 *, 

316 compression_options: Mapping[str, Any] | None = None, 

317 opaque_metadata: FitsOpaqueMetadata | None = None, 

318 root: Ndf | NdfContainer | None = None, 

319 lsst_path: str = "/MORE/LSST", 

320 direct_ndf_array_paths: Mapping[str, str] | None = None, 

321 wcs_ndf_paths: Sequence[str] = ("/",), 

322 ) -> Iterator[Self]: 

323 """Open an NDF file for writing and yield an `NdfOutputArchive`. 

324 

325 ``filename=None`` uses an in-memory HDF5 file; the on-disk 

326 artefact is discarded but the archive's writes still produce a 

327 usable returned tree (handy for tests). 

328 """ 

329 if filename is None: 329 ↛ 330line 329 didn't jump to line 330 because the condition on line 329 was never true

330 h5_file = h5py.File("inmem.sdf", "w", driver="core", backing_store=False) 

331 else: 

332 if os.path.exists(filename) and os.path.getsize(filename) > 0: 332 ↛ 333line 332 didn't jump to line 333 because the condition on line 332 was never true

333 raise OSError(f"File {filename!r} already exists.") 

334 h5_file = h5py.File(filename, "w") 

335 try: 

336 yield cls( 

337 h5_file, 

338 compression_options=compression_options, 

339 opaque_metadata=opaque_metadata, 

340 root=root, 

341 lsst_path=lsst_path, 

342 direct_ndf_array_paths=direct_ndf_array_paths, 

343 wcs_ndf_paths=wcs_ndf_paths, 

344 ) 

345 finally: 

346 h5_file.close() 

347 

348 def add_tree( 

349 self, 

350 tree: ArchiveTree, 

351 *, 

352 sky_projection: Any = None, 

353 bbox: Any = None, 

354 unit: astropy.units.UnitBase | None = None, 

355 root_name: str | None = None, 

356 ) -> None: 

357 """Finalize the file: write WCS, units, JSON tree, and ORIGIN. 

358 

359 Writes the canonical NDF ``/WCS`` HDS structure (an AST channel 

360 text dump that KAPPA / hdstrace expect) when ``sky_projection`` is 

361 provided. A native mask sub-NDF at ``/MORE/LSST/MASK`` gets a 3D 

362 WCS whose first two axes reuse the parent's sky projection and 

363 whose third axis is the mask-byte coordinate. The JSON tree at 

364 ``<lsst_path>/JSON`` remains the source of truth for symmetric 

365 round-trips; ``/WCS`` is for Starlink tools. Auto-detect read of 

366 ``/WCS/DATA`` into a typed ``SkyProjection`` is a follow-up. 

367 

368 Parameters 

369 ---------- 

370 tree 

371 Pydantic tree returned by the object's ``serialize`` method, 

372 with ``metadata``/``butler_info`` already applied. 

373 sky_projection, bbox, unit 

374 Top-level object attributes that drive NDF-canonical writes. 

375 root_name 

376 Value to assign to the root group's ``HDS_ROOT_NAME`` 

377 attribute (fixed-length ASCII so KAPPA / hdstrace decode it). 

378 """ 

379 if sky_projection is not None: 

380 self._write_wcs(sky_projection, bbox) 

381 if unit is not None and isinstance(self._document.root, Ndf): 

382 self._document.ensure_ndf("/").set_units(_unit_to_ndf_string(unit)) 

383 json_text = tree.model_dump_json() 

384 # Let ensure_structure derive types from path-component names so 

385 # /MORE/LSST/... gets <MORE> stay as EXT and <LSST> typed LSST, 

386 # mirroring HDS convention. 

387 lsst = self._ensure_model_structure(self._lsst_path) 

388 lsst.children["JSON"] = HdsPrimitive.char_array([json_text], width=max(80, len(json_text))) 

389 lsst.children["DATA_MODEL"] = HdsPrimitive.char_scalar( 

390 tree.schema_url, width=max(80, len(tree.schema_url)) 

391 ) 

392 lsst.children["FORMAT_VERSION"] = HdsPrimitive.array(np.array(_NDF_FORMAT_VERSION, dtype=np.int32)) 

393 primary = self._opaque_metadata.headers.get(ExtensionKey()) 

394 if primary is not None and len(primary): 

395 cards = _fits_header_records(primary) 

396 more = self._ensure_model_structure("/MORE") 

397 more.children["FITS"] = HdsPrimitive.char_array(cards, width=80) 

398 if bbox is not None: 398 ↛ 403line 398 didn't jump to line 403 because the condition on line 398 was always true

399 origin = _origin_from_bbox(bbox) 

400 for struct_path in self._bbox_array_struct_paths: 

401 if self._has_model_path(struct_path): 401 ↛ 400line 401 didn't jump to line 400 because the condition on line 401 was always true

402 self.set_array_origin(struct_path, origin) 

403 if root_name is not None: 403 ↛ 405line 403 didn't jump to line 405 because the condition on line 403 was always true

404 self._document.root_name = root_name 

405 self._flush() 

406 

407 def _write_wcs(self, sky_projection: Any, bbox: Any) -> None: 

408 ast_frame_set = sky_projection._pixel_to_sky._get_ast_frame_set() 

409 text = _show_ast_for_ndf(ast_frame_set, bbox) 

410 lines = _hds.encode_ndf_ast_data(text) 

411 for ndf_path in self._wcs_ndf_paths: 

412 if self._has_model_path(ndf_path): 412 ↛ 411line 412 didn't jump to line 411 because the condition on line 412 was always true

413 self._document.ensure_ndf(ndf_path).set_wcs(NdfWcs(lines)) 

414 if self._has_model_path("/MORE/LSST/MASK"): 

415 mask_origin = _origin_from_bbox(bbox) if bbox is not None else (0, 0) 

416 mask_ndim = self._model_array_ndim("/MORE/LSST/MASK/DATA_ARRAY") 

417 mask_origin = (*mask_origin, *((0,) * max(0, mask_ndim - len(mask_origin)))) 

418 mask_ast_frame_set = sky_projection._pixel_to_sky._get_ast_frame_set() 

419 mask_text = _show_mask_ast_for_ndf( 

420 mask_ast_frame_set, 

421 mask_origin, 

422 labels=("x", "y", "mask-byte"), 

423 ) 

424 self._document.ensure_ndf("/MORE/LSST/MASK").set_wcs(NdfWcs(_hds.encode_ndf_ast_data(mask_text))) 

425 

426 def serialize_direct[T: pydantic.BaseModel | None]( 

427 self, name: str, serializer: Callable[[OutputArchive[NdfPointerModel]], T] 

428 ) -> T: 

429 nested = NestedOutputArchive[NdfPointerModel](name, self) 

430 return serializer(nested) 

431 

432 def serialize_pointer[T: ArchiveTree]( 

433 self, name: str, serializer: Callable[[OutputArchive[NdfPointerModel]], T], key: Hashable 

434 ) -> NdfPointerModel: 

435 if (pointer := self._pointers.get(key)) is not None: 

436 return pointer 

437 archive_path = name if name.startswith("/") else f"/{name}" 

438 target_path = self._archive_path_to_hdf5_path(archive_path) 

439 self._register_hdf5_path(target_path, archive_path) 

440 # Run the serializer first so any nested add_array / serialize_pointer 

441 # calls write into the file before we dump this sub-tree to JSON. 

442 model = self.serialize_direct(name, serializer) 

443 json_text = model.model_dump_json() 

444 # Store the JSON document as a "JSON" child of the target structure, 

445 # mirroring the main /MORE/LSST/JSON tree. Writing it at the target 

446 # path itself would clobber any nested arrays the serializer added 

447 # under that path (IR.write_to_hdf5 deletes the existing node before 

448 # writing a primitive). 

449 # ensure_structure derives the HDS type from the leaf name, so 

450 # /MORE/LSST/PSF is typed <PSF> rather than the generic <EXT>. 

451 target = self._ensure_model_structure(target_path) 

452 target.children["JSON"] = HdsPrimitive.char_array([json_text], width=max(80, len(json_text))) 

453 self._flush() 

454 pointer = NdfPointerModel(path=f"{target_path}/JSON") 

455 self._pointers[key] = pointer 

456 return pointer 

457 

458 def serialize_frame_set[T: ArchiveTree]( 

459 self, name: str, frame_set: FrameSet, serializer: Callable[[OutputArchive], T], key: Hashable 

460 ) -> NdfPointerModel: 

461 pointer = self.serialize_pointer(name, serializer, key) 

462 self._frame_sets.append((frame_set, pointer)) 

463 return pointer 

464 

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

466 return iter(self._frame_sets) 

467 

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

469 _prefer_native_mask_arrays = True 

470 

471 def add_array( 

472 self, 

473 array: np.ndarray, 

474 *, 

475 name: str | None = None, 

476 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates, 

477 tile_shape: tuple[int, ...] | None = None, 

478 options_name: str | None = None, 

479 ) -> ArrayReferenceModel: 

480 if name is None: 480 ↛ 481line 480 didn't jump to line 481 because the condition on line 480 was never true

481 raise ValueError("Anonymous arrays are not supported in the NDF archive.") 

482 name, version = self._register_name(name) 

483 # Recognised top-level names go to standard NDF locations. 

484 # Anything else hoists under /MORE/LSST. 

485 if name == "image": 

486 root = self._document.ensure_ndf("/") 

487 root.set_array_component( 

488 "DATA_ARRAY", 

489 array, 

490 origin=np.zeros(array.ndim, dtype=np.int64), 

491 compression_options=self._compression_options, 

492 ) 

493 path = "/DATA_ARRAY/DATA" 

494 self._bbox_array_struct_paths.add("/DATA_ARRAY") 

495 elif name == "variance": 

496 root = self._document.ensure_ndf("/") 

497 root.set_array_component( 

498 "VARIANCE", 

499 array, 

500 origin=np.zeros(array.ndim, dtype=np.int64), 

501 compression_options=self._compression_options, 

502 ) 

503 path = "/VARIANCE/DATA" 

504 self._bbox_array_struct_paths.add("/VARIANCE") 

505 elif name == "mask": 

506 if array.ndim == 2 and array.dtype in self._COMPATIBLE_MASK_DTYPES: 

507 self._set_quality_array(array) 

508 path = "/QUALITY/QUALITY/DATA" 

509 self._bbox_array_struct_paths.add("/QUALITY/QUALITY") 

510 else: 

511 # Native Mask serialization passes HDF5 shape 

512 # (mask-byte, y, x). HDS reports the reverse dimension order, 

513 # so Starlink tools see (x, y, mask-byte). 

514 if array.ndim == 3 and array.dtype in self._COMPATIBLE_MASK_DTYPES: 514 ↛ 517line 514 didn't jump to line 517 because the condition on line 514 was always true

515 self._set_quality_array(self._collapse_mask_to_quality(array)) 

516 self._bbox_array_struct_paths.add("/QUALITY/QUALITY") 

517 mask_ndf = self._document.ensure_ndf("/MORE/LSST/MASK") 

518 mask_ndf.set_array_component( 

519 "DATA_ARRAY", 

520 array, 

521 origin=np.zeros(array.ndim, dtype=np.int64), 

522 bad_pixel=False, 

523 compression_options=self._compression_options, 

524 ) 

525 path = "/MORE/LSST/MASK/DATA_ARRAY/DATA" 

526 self._bbox_array_struct_paths.add("/MORE/LSST/MASK/DATA_ARRAY") 

527 elif name in self._direct_ndf_array_paths: 

528 sub_ndf_path = self._direct_ndf_array_paths[name] 

529 sub_ndf = self._document.ensure_ndf(sub_ndf_path) 

530 sub_ndf.set_array_component( 

531 "DATA_ARRAY", 

532 array, 

533 origin=np.zeros(array.ndim, dtype=np.int64), 

534 compression_options=self._compression_options, 

535 ) 

536 path = f"{sub_ndf_path}/DATA_ARRAY/DATA" 

537 self._bbox_array_struct_paths.add(f"{sub_ndf_path}/DATA_ARRAY") 

538 else: 

539 # Hoisted numeric arrays are wrapped as sub-NDFs under 

540 # /MORE/LSST/<UPPER_PATH> so Starlink tools (KAPPA `display`, 

541 # `hdstrace`, etc.) can inspect them just like the main image. 

542 # The sub-NDF has the canonical layout: top-level group with 

543 # CLASS="NDF" containing a DATA_ARRAY structure (CLASS="ARRAY") 

544 # with DATA + ORIGIN primitives. Over-long components are shrunk 

545 # to fit the HDS object-name limit (DAT__SZNAM); repeated names 

546 # get a version suffix on their leaf so siblings stay distinct. 

547 archive_path, logical_id = self._versioned_archive_path(name, version) 

548 sub_ndf_path = self._archive_path_to_hdf5_path(archive_path) 

549 self._register_hdf5_path(sub_ndf_path, logical_id) 

550 sub_ndf = self._document.ensure_ndf(sub_ndf_path) 

551 sub_ndf.set_array_component( 

552 "DATA_ARRAY", 

553 array, 

554 origin=np.zeros(array.ndim, dtype=np.int64), 

555 compression_options=self._compression_options, 

556 ) 

557 path = f"{sub_ndf_path}/DATA_ARRAY/DATA" 

558 self._flush() 

559 # Shape is stored in the JSON tree (matching the FITS archive) because 

560 # MaskSerializationModel.bbox needs it before any arrays are read. 

561 # Future work: resolve shape from the HDF5 dataset on read instead. 

562 return ArrayReferenceModel( 

563 source=f"ndf:{path}", 

564 shape=list(array.shape), 

565 datatype=NumberType.from_numpy(array.dtype), 

566 ) 

567 

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

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

570 

571 Intermediate groups created on demand are tagged with 

572 ``CLASS="EXT"``, the HDS type for general-purpose extension 

573 containers (matches what hds-v5 writes for ``MORE``). 

574 """ 

575 self._ensure_model_structure(path, "EXT") 

576 self._flush() 

577 return self._file["/" if path in ("", "/") else path] 

578 

579 def _ensure_struct(self, path: str, hds_type: str) -> h5py.Group: 

580 """Ensure a structure exists at ``path`` with the given HDS type.""" 

581 self._ensure_model_structure(path, hds_type) 

582 self._flush() 

583 return self._file["/" if path in ("", "/") else path] 

584 

585 def _ensure_array_structure(self, path: str) -> h5py.Group: 

586 """Ensure an HDS ``ARRAY`` structure exists at ``path``.""" 

587 return self._ensure_struct(path, "ARRAY") 

588 

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

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

591 

592 BADBITS is set to 255 so every bit of the collapsed single-byte 

593 QUALITY plane is treated as bad by NDF applications. 

594 """ 

595 root = self._document.ensure_ndf("/") 

596 if "QUALITY" not in root.children: 

597 root.children["QUALITY"] = HdsStructure("QUALITY") 

598 quality = root.get_structure("QUALITY") 

599 quality.hds_type = "QUALITY" 

600 quality.children["BADBITS"] = HdsPrimitive.array(np.array(255, dtype=np.uint8)) 

601 self._flush() 

602 return self._file["/QUALITY"] 

603 

604 def _ensure_quality_array_structure(self) -> h5py.Group: 

605 """Ensure the nested ``/QUALITY/QUALITY`` ARRAY structure exists.""" 

606 root = self._document.ensure_ndf("/") 

607 if "QUALITY" not in root.children: 

608 root.children["QUALITY"] = HdsStructure("QUALITY") 

609 quality = root.get_structure("QUALITY") 

610 quality.hds_type = "QUALITY" 

611 quality.children["BADBITS"] = HdsPrimitive.array(np.array(255, dtype=np.uint8)) 

612 if "QUALITY" not in quality.children or not isinstance(quality.children["QUALITY"], HdsStructure): 

613 quality.children["QUALITY"] = HdsStructure("ARRAY") 

614 quality_array = quality.get_structure("QUALITY") 

615 quality_array.hds_type = "ARRAY" 

616 quality_array.children.setdefault("ORIGIN", HdsPrimitive.array(np.zeros(2, dtype=np.int32))) 

617 quality_array.children.setdefault("BAD_PIXEL", HdsPrimitive.array(np.array(False, dtype=np.bool_))) 

618 self._flush() 

619 return self._file["/QUALITY/QUALITY"] 

620 

621 def _write_quality_array(self, quality: np.ndarray) -> None: 

622 """Write or replace the NDF QUALITY array.""" 

623 self._set_quality_array(quality) 

624 self._flush() 

625 

626 def _set_quality_array(self, quality: np.ndarray) -> None: 

627 """Set or replace the NDF QUALITY array in the model.""" 

628 root = self._document.ensure_ndf("/") 

629 root.set_quality( 

630 NdfQuality( 

631 NdfArray( 

632 quality, 

633 origin=np.zeros(2, dtype=np.int32), 

634 bad_pixel=False, 

635 compression_options=self._compression_options, 

636 ) 

637 ) 

638 ) 

639 

640 def _collapse_mask_to_quality(self, array: np.ndarray) -> np.ndarray: 

641 """Compress an NDF-native 3-D mask array into 2-D QUALITY. 

642 

643 The input array is in HDF5 storage order ``(mask-byte, y, x)``. 

644 Single-byte masks copy directly to preserve bit values. Wider masks 

645 collapse to 1 where any byte is non-zero and 0 otherwise. 

646 """ 

647 if array.shape[0] == 1: 

648 return array[0, :, :] 

649 return np.any(array != 0, axis=0).astype(np.uint8) 

650 

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

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

653 

654 The top-level `write` function overwrites this 

655 with bbox-derived values via :meth:`set_array_origin` once the 

656 bbox is known. 

657 """ 

658 struct = self._document.root.get_structure(struct_path) 

659 if "ORIGIN" not in struct.children: 

660 struct.children["ORIGIN"] = HdsPrimitive.array(np.zeros(array.ndim, dtype=np.int64)) 

661 

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

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

664 

665 Parameters 

666 ---------- 

667 struct_path 

668 HDF5 path to the ARRAY structure (e.g. ``"/DATA_ARRAY"``). 

669 origin 

670 Origin in NDF/Fortran axis order (e.g. ``(x_min, y_min)`` 

671 for a 2D image with bbox lower bound ``(x_min, y_min)``). 

672 """ 

673 struct = self._document.root.get_structure(struct_path) 

674 origin_dtype = np.int32 if struct_path == "/QUALITY/QUALITY" else np.int64 

675 origin_array = np.asarray(origin, dtype=origin_dtype) 

676 data_node = struct.children.get("DATA") 

677 if isinstance(data_node, HdsPrimitive): 677 ↛ 681line 677 didn't jump to line 681 because the condition on line 677 was always true

678 data_ndim = data_node.read_array().ndim 

679 if origin_array.size < data_ndim: 

680 origin_array = np.pad(origin_array, (0, data_ndim - origin_array.size)) 

681 struct.children["ORIGIN"] = HdsPrimitive.array(origin_array) 

682 self._flush() 

683 

684 def _ensure_model_structure(self, path: str, hds_type: str | None = None) -> HdsStructure: 

685 """Return or create a structure in the NDF document model. 

686 

687 With ``hds_type=None`` the leaf type is derived from its name; 

688 see `HdsStructure.ensure_structure`. 

689 """ 

690 if path in ("", "/"): 690 ↛ 691line 690 didn't jump to line 691 because the condition on line 690 was never true

691 return self._document.root 

692 return self._document.root.ensure_structure(path, hds_type) 

693 

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

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

696 if self._lsst_path == "/MORE/LSST": 696 ↛ 698line 696 didn't jump to line 698 because the condition on line 696 was always true

697 return archive_path_to_hdf5_path(archive_path, self._name_shrinker) 

698 if not archive_path: 

699 return f"{self._lsst_path}/JSON" 

700 components = archive_path_to_hdf5_path_components(archive_path, self._name_shrinker) 

701 return f"{self._lsst_path}/{'/'.join(components)}" 

702 

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

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

705 

706 ``logical_id`` is the un-shrunk archive path (with any version suffix 

707 applied), which is unique per logical write. 

708 Two different archive entries shrinking to the same HDS path would 

709 silently clobber one another, so this fails loudly instead. 

710 """ 

711 previous = self._hdf5_path_owners.get(hdf5_path) 

712 if previous is not None and previous != logical_id: 

713 raise ValueError( 

714 f"NDF/HDS name collision: archive entries {previous!r} and {logical_id!r} " 

715 f"both map to {hdf5_path!r} after shrinking to the {_hds.DAT__SZNAM}-character " 

716 f"HDS name limit; rename one of them." 

717 ) 

718 self._hdf5_path_owners[hdf5_path] = logical_id 

719 

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

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

722 

723 ``archive_path`` has any version suffix applied to its leaf through the 

724 version-aware shrinker (so the suffix survives the later per-component 

725 shrink); ``logical_id`` is the un-shrunk version-applied path used for 

726 collision detection. 

727 """ 

728 archive_path = name if name.startswith("/") else f"/{name}" 

729 if version <= 1: 

730 return archive_path, archive_path 

731 head, sep, leaf = archive_path.rpartition("/") 

732 logical_id = f"{archive_path}_{version}" 

733 versioned = f"{head}{sep}{self._name_shrinker.shrink_versioned(leaf, version)}" 

734 return versioned, logical_id 

735 

736 def _has_model_path(self, path: str) -> bool: 

737 """Return `True` if a path exists in the NDF document model.""" 

738 try: 

739 self._document.get(path) 

740 except KeyError: 

741 return False 

742 return True 

743 

744 def _model_array_ndim(self, struct_path: str) -> int: 

745 """Return the dimensionality of an ARRAY structure's DATA primitive.""" 

746 struct = self._document.root.get_structure(struct_path) 

747 data_node = struct.children["DATA"] 

748 if not isinstance(data_node, HdsPrimitive): 748 ↛ 749line 748 didn't jump to line 749 because the condition on line 748 was never true

749 raise TypeError(f"{struct_path}/DATA is not an HDS primitive.") 

750 return data_node.read_array().ndim 

751 

752 def _flush(self) -> None: 

753 """Synchronize the Python NDF document model to the open HDF5 file.""" 

754 self._document.write_to_hdf5(self._file) 

755 

756 def add_table( 

757 self, 

758 table: astropy.table.Table, 

759 *, 

760 name: str | None = None, 

761 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates, 

762 ) -> TableModel: