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 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 if archive_default_name is not None: 

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

117 else: 

118 tree = obj.serialize(archive) 

119 if metadata is not None: 

120 tree.metadata.update(metadata) 

121 if butler_info is not None: 

122 tree.butler_info = butler_info 

123 archive.add_tree( 

124 tree, 

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

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

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

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

129 ) 

130 return tree 

131 

132 

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

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

135 

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

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

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

139 ``(x_min, y_min)``. 

140 """ 

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

142 # .start attribute (the lower bound). 

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

144 x = bbox.x 

145 y = bbox.y 

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

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

148 raise AttributeError( 

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

150 "_origin_from_bbox needs updating." 

151 ) 

152 

153 

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

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

156 try: 

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

158 except ValueError: 

159 return unit.to_string() 

160 

161 

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

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

164 

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

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

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

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

169 header block. 

170 """ 

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

172 encoded = block.encode("ascii") 

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

174 raise ValueError( 

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

176 ) 

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

178 

179 

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

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

182 if isinstance(obj, ColorImage): 

183 return { 

184 "root": NdfContainer(), 

185 "lsst_path": "/LSST", 

186 "direct_ndf_array_paths": { 

187 "red": "/RED", 

188 "green": "/GREEN", 

189 "blue": "/BLUE", 

190 }, 

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

192 } 

193 return {} 

194 

195 

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

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

198 

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

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

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

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

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

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

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

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

207 """ 

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

209 x_shift = 1.0 

210 y_shift = 1.0 

211 else: 

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

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

214 

215 saved_current = ast_frame_set.current 

216 ast_frame_set.current = ast_frame_set.base 

217 ast_frame_set.domain = "PIXEL" 

218 ast_frame_set.current = saved_current 

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

220 

221 stream = StringStream() 

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

223 channel.write(ast_frame_set) 

224 return stream.getSinkData() 

225 

226 

227def _show_mask_ast_for_ndf( 

228 parent_ast_frame_set: Any, 

229 origin: Sequence[int], 

230 *, 

231 labels: Sequence[str] = (), 

232) -> str: 

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

234 

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

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

237 """ 

238 n_axes = len(origin) 

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

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

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

242 pixel_frame.setLabel(axis, label) 

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

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

245 

246 parent_pixel_to_sky = parent_ast_frame_set.getMapping( 

247 parent_ast_frame_set.base, 

248 parent_ast_frame_set.current, 

249 ) 

250 parent_sky_frame = parent_ast_frame_set.getFrame(parent_ast_frame_set.current) 

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

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

253 ast_frame_set.addFrame( 

254 ast_frame_set.current, 

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

256 CmpFrame(parent_sky_frame, mask_axis_frame), 

257 ) 

258 

259 stream = StringStream() 

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

261 channel.write(ast_frame_set) 

262 return stream.getSinkData() 

263 

264 

265class NdfOutputArchive(OutputArchive[NdfPointerModel]): 

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

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

268 model. 

269 

270 Parameters 

271 ---------- 

272 file 

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

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

275 compression_options 

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

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

278 opaque_metadata 

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

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

281 top-level `write` function. 

282 root 

283 Root NDF or NDF container to populate; a new empty ``Ndf`` is created 

284 when `None`. 

285 lsst_path 

286 HDS path under which the LSST extension is written. 

287 direct_ndf_array_paths 

288 Mapping from archive path to NDF array path for arrays written 

289 directly as NDF components rather than into the LSST extension. 

290 wcs_ndf_paths 

291 NDF paths for which WCS information is written. 

292 """ 

293 

294 def __init__( 

295 self, 

296 file: h5py.File, 

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

298 opaque_metadata: FitsOpaqueMetadata | None = None, 

299 root: Ndf | NdfContainer | None = None, 

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

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

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

303 ) -> None: 

304 super().__init__() 

305 self._file = file 

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

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

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

309 self._wcs_ndf_paths = tuple(wcs_ndf_paths) 

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

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

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

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

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

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

316 self._name_shrinker = HdsNameShrinker() 

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

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

319 self._flush() 

320 

321 @classmethod 

322 @contextmanager 

323 def open( 

324 cls, 

325 filename: str | Path | None, 

326 *, 

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

328 opaque_metadata: FitsOpaqueMetadata | None = None, 

329 root: Ndf | NdfContainer | None = None, 

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

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

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

333 ) -> Iterator[Self]: 

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

335 

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

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

338 usable returned tree (handy for tests). 

339 

340 Parameters 

341 ---------- 

342 filename 

343 Name of the file to write to, or `None` to use an in-memory 

344 HDF5 file whose on-disk artefact is discarded. 

345 compression_options 

346 Optional dict passed through to `h5py.Group.create_dataset` 

347 for image arrays. 

348 opaque_metadata 

349 Optional `~lsst.images.fits.FitsOpaqueMetadata` whose cards are 

350 written to ``/MORE/FITS`` by the top-level `write` function. 

351 root 

352 Root NDF or NDF container to populate; a new empty ``Ndf`` is 

353 created when `None`. 

354 lsst_path 

355 HDS path under which the LSST extension is written. 

356 direct_ndf_array_paths 

357 Mapping from archive path to NDF array path for arrays written 

358 directly as NDF components rather than into the LSST extension. 

359 wcs_ndf_paths 

360 NDF paths for which WCS information is written. 

361 """ 

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

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

364 else: 

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

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

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

368 try: 

369 yield cls( 

370 h5_file, 

371 compression_options=compression_options, 

372 opaque_metadata=opaque_metadata, 

373 root=root, 

374 lsst_path=lsst_path, 

375 direct_ndf_array_paths=direct_ndf_array_paths, 

376 wcs_ndf_paths=wcs_ndf_paths, 

377 ) 

378 finally: 

379 h5_file.close() 

380 

381 def add_tree( 

382 self, 

383 tree: ArchiveTree, 

384 *, 

385 sky_projection: Any = None, 

386 bbox: Any = None, 

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

388 root_name: str | None = None, 

389 ) -> None: 

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

391 

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

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

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

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

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

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

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

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

400 

401 Parameters 

402 ---------- 

403 tree 

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

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

406 sky_projection, bbox, unit 

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

408 root_name 

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

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

411 """ 

412 if sky_projection is not None: 

413 self._write_wcs(sky_projection, bbox) 

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

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

416 json_text = tree.model_dump_json() 

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

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

419 # mirroring HDS convention. 

420 lsst = self._ensure_model_structure(self._lsst_path) 

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

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

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

424 ) 

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

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

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

428 cards = _fits_header_records(primary) 

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

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

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

432 origin = _origin_from_bbox(bbox) 

433 for struct_path in self._bbox_array_struct_paths: 

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

435 self.set_array_origin(struct_path, origin) 

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

437 self._document.root_name = root_name 

438 self._flush() 

439 

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

441 ast_frame_set = sky_projection._pixel_to_sky._get_ast_frame_set() 

442 text = _show_ast_for_ndf(ast_frame_set, bbox) 

443 lines = _hds.encode_ndf_ast_data(text) 

444 for ndf_path in self._wcs_ndf_paths: 

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

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

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

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

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

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

451 mask_ast_frame_set = sky_projection._pixel_to_sky._get_ast_frame_set() 

452 mask_text = _show_mask_ast_for_ndf( 

453 mask_ast_frame_set, 

454 mask_origin, 

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

456 ) 

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

458 

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

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

461 ) -> T: 

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

463 return serializer(nested) 

464 

465 def serialize_pointer[T: ArchiveTree]( 

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

467 ) -> NdfPointerModel: 

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

469 return pointer 

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

471 target_path = self._archive_path_to_hdf5_path(archive_path) 

472 self._register_hdf5_path(target_path, archive_path) 

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

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

475 model = self.serialize_direct(name, serializer) 

476 json_text = model.model_dump_json() 

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

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

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

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

481 # writing a primitive). 

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

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

484 target = self._ensure_model_structure(target_path) 

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

486 self._flush() 

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

488 self._pointers[key] = pointer 

489 return pointer 

490 

491 def serialize_frame_set[T: ArchiveTree]( 

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

493 ) -> NdfPointerModel: 

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

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

496 return pointer 

497 

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

499 return iter(self._frame_sets) 

500 

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

502 _prefer_native_mask_arrays = True 

503 

504 def add_array( 

505 self, 

506 array: np.ndarray, 

507 *, 

508 name: str | None = None, 

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

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

511 options_name: str | None = None, 

512 ) -> ArrayReferenceModel: 

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

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

515 name, version = self._register_name(name) 

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

517 # Anything else hoists under /MORE/LSST. 

518 if name == "image": 

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

520 root.set_array_component( 

521 "DATA_ARRAY", 

522 array, 

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

524 compression_options=self._compression_options, 

525 ) 

526 path = "/DATA_ARRAY/DATA" 

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

528 elif name == "variance": 

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

530 root.set_array_component( 

531 "VARIANCE", 

532 array, 

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

534 compression_options=self._compression_options, 

535 ) 

536 path = "/VARIANCE/DATA" 

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

538 elif name == "mask": 

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

540 self._set_quality_array(array) 

541 path = "/QUALITY/QUALITY/DATA" 

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

543 else: 

544 # Native Mask serialization passes HDF5 shape 

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

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

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

548 self._set_quality_array(self._collapse_mask_to_quality(array)) 

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

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

551 mask_ndf.set_array_component( 

552 "DATA_ARRAY", 

553 array, 

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

555 bad_pixel=False, 

556 compression_options=self._compression_options, 

557 ) 

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

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

560 elif name in self._direct_ndf_array_paths: 

561 sub_ndf_path = self._direct_ndf_array_paths[name] 

562 sub_ndf = self._document.ensure_ndf(sub_ndf_path) 

563 sub_ndf.set_array_component( 

564 "DATA_ARRAY", 

565 array, 

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

567 compression_options=self._compression_options, 

568 ) 

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

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

571 else: 

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

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

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

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

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

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

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

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

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

581 sub_ndf_path = self._archive_path_to_hdf5_path(archive_path) 

582 self._register_hdf5_path(sub_ndf_path, logical_id) 

583 sub_ndf = self._document.ensure_ndf(sub_ndf_path) 

584 sub_ndf.set_array_component( 

585 "DATA_ARRAY", 

586 array, 

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

588 compression_options=self._compression_options, 

589 ) 

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

591 self._flush() 

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

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

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

595 return ArrayReferenceModel( 

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

597 shape=list(array.shape), 

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

599 ) 

600 

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

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

603 

604 Intermediate groups created on demand are tagged with 

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

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

607 """ 

608 self._ensure_model_structure(path, "EXT") 

609 self._flush() 

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

611 

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

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

614 self._ensure_model_structure(path, hds_type) 

615 self._flush() 

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

617 

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

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

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

621 

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

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

624 

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

626 QUALITY plane is treated as bad by NDF applications. 

627 """ 

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

629 if "QUALITY" not in root.children: 

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

631 quality = root.get_structure("QUALITY") 

632 quality.hds_type = "QUALITY" 

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

634 self._flush() 

635 return self._file["/QUALITY"] 

636 

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

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

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

640 if "QUALITY" not in root.children: 

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

642 quality = root.get_structure("QUALITY") 

643 quality.hds_type = "QUALITY" 

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

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

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

647 quality_array = quality.get_structure("QUALITY") 

648 quality_array.hds_type = "ARRAY" 

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

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

651 self._flush() 

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

653 

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

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

656 self._set_quality_array(quality) 

657 self._flush() 

658 

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

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

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

662 root.set_quality( 

663 NdfQuality( 

664 NdfArray( 

665 quality, 

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

667 bad_pixel=False, 

668 compression_options=self._compression_options, 

669 ) 

670 ) 

671 ) 

672 

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

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

675 

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

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

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

679 """ 

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

681 return array[0, :, :] 

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

683 

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

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

686 

687 The top-level `write` function overwrites this 

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

689 bbox is known. 

690 """ 

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

692 if "ORIGIN" not in struct.children: 

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

694 

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

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

697 

698 Parameters 

699 ---------- 

700 struct_path 

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

702 origin 

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

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

705 """ 

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

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

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

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

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

711 data_ndim = data_node.read_array().ndim 

712 if origin_array.size < data_ndim: 

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

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

715 self._flush() 

716 

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

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

719 

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

721 see `HdsStructure.ensure_structure`. 

722 """ 

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

724 return self._document.root 

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

726 

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

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

729 

730 Parameters 

731 ---------- 

732 archive_path 

733 Serialization archive path to translate. 

734 """ 

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

736 return archive_path_to_hdf5_path(archive_path, self._name_shrinker) 

737 if not archive_path: 

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

739 components = archive_path_to_hdf5_path_components(archive_path, self._name_shrinker) 

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

741 

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

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

744 

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

746 applied), which is unique per logical write. 

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

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

749 

750 Parameters 

751 ---------- 

752 hdf5_path 

753 HDF5 path being claimed. 

754 logical_id 

755 Un-shrunk archive path that owns ``hdf5_path``. 

756 """ 

757 previous = self._hdf5_path_owners.get(hdf5_path) 

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

759 raise ValueError( 

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

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

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

763 ) 

764 self._hdf5_path_owners[hdf5_path] = logical_id 

765 

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

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

768 

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

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

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

772 collision detection. 

773 """ 

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

775 if version <= 1: 

776 return archive_path, archive_path 

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

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

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

780 return versioned, logical_id 

781 

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