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 root 

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

283 when `None`. 

284 lsst_path 

285 HDS path under which the LSST extension is written. 

286 direct_ndf_array_paths 

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

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

289 wcs_ndf_paths 

290 NDF paths for which WCS information is written. 

291 """ 

292 

293 def __init__( 

294 self, 

295 file: h5py.File, 

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

297 opaque_metadata: FitsOpaqueMetadata | None = None, 

298 root: Ndf | NdfContainer | None = None, 

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

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

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

302 ) -> None: 

303 super().__init__() 

304 self._file = file 

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

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

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

308 self._wcs_ndf_paths = tuple(wcs_ndf_paths) 

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

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

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

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

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

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

315 self._name_shrinker = HdsNameShrinker() 

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

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

318 self._flush() 

319 

320 @classmethod 

321 @contextmanager 

322 def open( 

323 cls, 

324 filename: str | None, 

325 *, 

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

327 opaque_metadata: FitsOpaqueMetadata | None = None, 

328 root: Ndf | NdfContainer | None = None, 

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

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

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

332 ) -> Iterator[Self]: 

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

334 

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

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

337 usable returned tree (handy for tests). 

338 

339 Parameters 

340 ---------- 

341 filename 

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

343 HDF5 file whose on-disk artefact is discarded. 

344 compression_options 

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

346 for image arrays. 

347 opaque_metadata 

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

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

350 root 

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

352 created when `None`. 

353 lsst_path 

354 HDS path under which the LSST extension is written. 

355 direct_ndf_array_paths 

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

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

358 wcs_ndf_paths 

359 NDF paths for which WCS information is written. 

360 """ 

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

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

363 else: 

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

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

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

367 try: 

368 yield cls( 

369 h5_file, 

370 compression_options=compression_options, 

371 opaque_metadata=opaque_metadata, 

372 root=root, 

373 lsst_path=lsst_path, 

374 direct_ndf_array_paths=direct_ndf_array_paths, 

375 wcs_ndf_paths=wcs_ndf_paths, 

376 ) 

377 finally: 

378 h5_file.close() 

379 

380 def add_tree( 

381 self, 

382 tree: ArchiveTree, 

383 *, 

384 sky_projection: Any = None, 

385 bbox: Any = None, 

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

387 root_name: str | None = None, 

388 ) -> None: 

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

390 

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

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

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

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

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

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

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

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

399 

400 Parameters 

401 ---------- 

402 tree 

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

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

405 sky_projection, bbox, unit 

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

407 root_name 

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

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

410 """ 

411 if sky_projection is not None: 

412 self._write_wcs(sky_projection, bbox) 

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

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

415 json_text = tree.model_dump_json() 

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

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

418 # mirroring HDS convention. 

419 lsst = self._ensure_model_structure(self._lsst_path) 

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

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

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

423 ) 

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

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

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

427 cards = _fits_header_records(primary) 

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

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

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

431 origin = _origin_from_bbox(bbox) 

432 for struct_path in self._bbox_array_struct_paths: 

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

434 self.set_array_origin(struct_path, origin) 

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

436 self._document.root_name = root_name 

437 self._flush() 

438 

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

440 ast_frame_set = sky_projection._pixel_to_sky._get_ast_frame_set() 

441 text = _show_ast_for_ndf(ast_frame_set, bbox) 

442 lines = _hds.encode_ndf_ast_data(text) 

443 for ndf_path in self._wcs_ndf_paths: 

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

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

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

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

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

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

450 mask_ast_frame_set = sky_projection._pixel_to_sky._get_ast_frame_set() 

451 mask_text = _show_mask_ast_for_ndf( 

452 mask_ast_frame_set, 

453 mask_origin, 

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

455 ) 

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

457 

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

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

460 ) -> T: 

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

462 return serializer(nested) 

463 

464 def serialize_pointer[T: ArchiveTree]( 

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

466 ) -> NdfPointerModel: 

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

468 return pointer 

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

470 target_path = self._archive_path_to_hdf5_path(archive_path) 

471 self._register_hdf5_path(target_path, archive_path) 

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

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

474 model = self.serialize_direct(name, serializer) 

475 json_text = model.model_dump_json() 

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

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

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

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

480 # writing a primitive). 

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

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

483 target = self._ensure_model_structure(target_path) 

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

485 self._flush() 

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

487 self._pointers[key] = pointer 

488 return pointer 

489 

490 def serialize_frame_set[T: ArchiveTree]( 

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

492 ) -> NdfPointerModel: 

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

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

495 return pointer 

496 

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

498 return iter(self._frame_sets) 

499 

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

501 _prefer_native_mask_arrays = True 

502 

503 def add_array( 

504 self, 

505 array: np.ndarray, 

506 *, 

507 name: str | None = None, 

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

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

510 options_name: str | None = None, 

511 ) -> ArrayReferenceModel: 

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

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

514 name, version = self._register_name(name) 

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

516 # Anything else hoists under /MORE/LSST. 

517 if name == "image": 

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

519 root.set_array_component( 

520 "DATA_ARRAY", 

521 array, 

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

523 compression_options=self._compression_options, 

524 ) 

525 path = "/DATA_ARRAY/DATA" 

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

527 elif name == "variance": 

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

529 root.set_array_component( 

530 "VARIANCE", 

531 array, 

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

533 compression_options=self._compression_options, 

534 ) 

535 path = "/VARIANCE/DATA" 

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

537 elif name == "mask": 

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

539 self._set_quality_array(array) 

540 path = "/QUALITY/QUALITY/DATA" 

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

542 else: 

543 # Native Mask serialization passes HDF5 shape 

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

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

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

547 self._set_quality_array(self._collapse_mask_to_quality(array)) 

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

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

550 mask_ndf.set_array_component( 

551 "DATA_ARRAY", 

552 array, 

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

554 bad_pixel=False, 

555 compression_options=self._compression_options, 

556 ) 

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

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

559 elif name in self._direct_ndf_array_paths: 

560 sub_ndf_path = self._direct_ndf_array_paths[name] 

561 sub_ndf = self._document.ensure_ndf(sub_ndf_path) 

562 sub_ndf.set_array_component( 

563 "DATA_ARRAY", 

564 array, 

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

566 compression_options=self._compression_options, 

567 ) 

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

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

570 else: 

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

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

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

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

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

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

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

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

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

580 sub_ndf_path = self._archive_path_to_hdf5_path(archive_path) 

581 self._register_hdf5_path(sub_ndf_path, logical_id) 

582 sub_ndf = self._document.ensure_ndf(sub_ndf_path) 

583 sub_ndf.set_array_component( 

584 "DATA_ARRAY", 

585 array, 

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

587 compression_options=self._compression_options, 

588 ) 

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

590 self._flush() 

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

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

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

594 return ArrayReferenceModel( 

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

596 shape=list(array.shape), 

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

598 ) 

599 

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

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

602 

603 Intermediate groups created on demand are tagged with 

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

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

606 """ 

607 self._ensure_model_structure(path, "EXT") 

608 self._flush() 

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

610 

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

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

613 self._ensure_model_structure(path, hds_type) 

614 self._flush() 

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

616 

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

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

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

620 

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

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

623 

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

625 QUALITY plane is treated as bad by NDF applications. 

626 """ 

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

628 if "QUALITY" not in root.children: 

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

630 quality = root.get_structure("QUALITY") 

631 quality.hds_type = "QUALITY" 

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

633 self._flush() 

634 return self._file["/QUALITY"] 

635 

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

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

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

639 if "QUALITY" not in root.children: 

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

641 quality = root.get_structure("QUALITY") 

642 quality.hds_type = "QUALITY" 

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

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

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

646 quality_array = quality.get_structure("QUALITY") 

647 quality_array.hds_type = "ARRAY" 

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

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

650 self._flush() 

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

652 

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

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

655 self._set_quality_array(quality) 

656 self._flush() 

657 

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

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

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

661 root.set_quality( 

662 NdfQuality( 

663 NdfArray( 

664 quality, 

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

666 bad_pixel=False, 

667 compression_options=self._compression_options, 

668 ) 

669 ) 

670 ) 

671 

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

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

674 

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

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

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

678 """ 

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

680 return array[0, :, :] 

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

682 

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

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

685 

686 The top-level `write` function overwrites this 

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

688 bbox is known. 

689 """ 

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

691 if "ORIGIN" not in struct.children: 

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

693 

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

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

696 

697 Parameters 

698 ---------- 

699 struct_path 

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

701 origin 

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

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

704 """ 

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

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

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

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

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

710 data_ndim = data_node.read_array().ndim 

711 if origin_array.size < data_ndim: 

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

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

714 self._flush() 

715 

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

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

718 

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

720 see `HdsStructure.ensure_structure`. 

721 """ 

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

723 return self._document.root 

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

725 

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

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

728 

729 Parameters 

730 ---------- 

731 archive_path 

732 Serialization archive path to translate. 

733 """ 

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

735 return archive_path_to_hdf5_path(archive_path, self._name_shrinker) 

736 if not archive_path: 

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

738 components = archive_path_to_hdf5_path_components(archive_path, self._name_shrinker) 

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

740 

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

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

743 

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

745 applied), which is unique per logical write. 

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

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

748 

749 Parameters 

750 ---------- 

751 hdf5_path 

752 HDF5 path being claimed. 

753 logical_id 

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

755 """ 

756 previous = self._hdf5_path_owners.get(hdf5_path) 

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

758 raise ValueError( 

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

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

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

762 ) 

763 self._hdf5_path_owners[hdf5_path] = logical_id 

764 

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

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

767 

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

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

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

771 collision detection. 

772 """ 

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

774 if version <= 1: 

775 return archive_path, archive_path 

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

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

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

779 return versioned, logical_id 

780 

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

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