Coverage for tests/test_ndf_input_archive.py: 99%

286 statements  

« prev     ^ index     » next       coverage.py v7.14.3, created at 2026-06-27 01:38 -0700

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 

14import os 

15import tempfile 

16import unittest 

17 

18import astropy.io.fits 

19import astropy.units as u 

20import numpy as np 

21import pydantic 

22 

23from lsst.images import Box, Image, ImageSerializationModel, Mask, MaskedImage 

24from lsst.images._transforms import FrameSet 

25from lsst.images.fits import ExtensionKey, FitsOpaqueMetadata 

26from lsst.images.serialization import ( 

27 ArchiveReadError, 

28 ArrayReferenceModel, 

29 InlineArrayModel, 

30 NumberType, 

31 read, 

32) 

33 

34try: 

35 import h5py 

36 

37 from lsst.images.ndf import ( 

38 NdfInputArchive, 

39 NdfOutputArchive, 

40 NdfPointerModel, 

41 _hds, 

42 read_starlink, 

43 write, 

44 ) 

45 

46 HAVE_H5PY = True 

47except ImportError: 

48 HAVE_H5PY = False 

49 

50 

51@unittest.skipUnless(HAVE_H5PY, "h5py is not installed") 

52class NdfInputArchiveOpenTestCase(unittest.TestCase): 

53 """Tests for `NdfInputArchive.open` and `get_tree`.""" 

54 

55 def test_open_round_trips_image_tree(self): 

56 image = Image( 

57 np.arange(20, dtype=np.float32).reshape(4, 5), 

58 bbox=Box.factory[10:14, 20:25], 

59 ) 

60 with tempfile.NamedTemporaryFile(suffix=".sdf", delete_on_close=False) as tmp: 

61 tmp.close() 

62 written_tree = write(image, tmp.name) 

63 with NdfInputArchive.open(tmp.name) as archive: 

64 tree = archive.get_tree(type(written_tree)) 

65 self.assertEqual(tree.model_dump_json(), written_tree.model_dump_json()) 

66 

67 def test_get_tree_raises_when_main_json_missing(self): 

68 # A file with no /MORE/LSST/JSON should raise ArchiveReadError. 

69 with tempfile.NamedTemporaryFile(suffix=".sdf", delete_on_close=False) as tmp: 

70 tmp.close() 

71 with h5py.File(tmp.name, "w") as f: 

72 f["/"].attrs["CLASS"] = "NDF" 

73 with NdfInputArchive.open(tmp.name) as archive: 

74 model_type = ImageSerializationModel[NdfPointerModel] 

75 with self.assertRaises(ArchiveReadError): 

76 archive.get_tree(model_type) 

77 

78 

79@unittest.skipUnless(HAVE_H5PY, "h5py is not installed") 

80class NdfInputArchiveDataTestCase(unittest.TestCase): 

81 """Tests for `get_array`, `deserialize_pointer`, and `get_frame_set`.""" 

82 

83 def test_get_array_reads_image_array(self): 

84 image = Image(np.arange(20, dtype=np.float32).reshape(4, 5)) 

85 with tempfile.NamedTemporaryFile(suffix=".sdf", delete_on_close=False) as tmp: 

86 tmp.close() 

87 tree = write(image, tmp.name) 

88 with NdfInputArchive.open(tmp.name) as archive: 

89 # The Image tree's `data` attribute is an 

90 # ArrayReferenceModel pointing at /DATA_ARRAY/DATA. 

91 arr = archive.get_array(tree.data) 

92 np.testing.assert_array_equal(arr, image.array) 

93 

94 def test_get_array_supports_slicing(self): 

95 image = Image(np.arange(20, dtype=np.float32).reshape(4, 5)) 

96 with tempfile.NamedTemporaryFile(suffix=".sdf", delete_on_close=False) as tmp: 

97 tmp.close() 

98 tree = write(image, tmp.name) 

99 with NdfInputArchive.open(tmp.name) as archive: 

100 arr = archive.get_array(tree.data, slices=(slice(0, 2), slice(1, 4))) 

101 np.testing.assert_array_equal(arr, image.array[:2, 1:4]) 

102 

103 def test_get_array_handles_inline_array(self): 

104 inline = InlineArrayModel(data=[1.0, 2.0, 3.0], datatype=NumberType.float64) 

105 image = Image(np.zeros((2, 2), dtype=np.float32)) 

106 with tempfile.NamedTemporaryFile(suffix=".sdf", delete_on_close=False) as tmp: 

107 tmp.close() 

108 write(image, tmp.name) 

109 with NdfInputArchive.open(tmp.name) as archive: 

110 arr = archive.get_array(inline) 

111 np.testing.assert_array_equal(arr, np.array([1.0, 2.0, 3.0])) 

112 

113 def test_get_array_unrecognised_source_raises(self): 

114 image = Image(np.zeros((2, 2), dtype=np.float32)) 

115 bogus = ArrayReferenceModel(source="fits:NOTUS", shape=[2, 2], datatype=NumberType.float32) 

116 with tempfile.NamedTemporaryFile(suffix=".sdf", delete_on_close=False) as tmp: 

117 tmp.close() 

118 write(image, tmp.name) 

119 with NdfInputArchive.open(tmp.name) as archive: 

120 with self.assertRaises(ArchiveReadError): 

121 archive.get_array(bogus) 

122 

123 def test_deserialize_pointer_round_trips_subtree(self): 

124 # Build a file with a hoisted sub-tree we can read back. Use the 

125 # output archive directly to avoid pulling in the full Image stack. 

126 class TinyTree(pydantic.BaseModel): 

127 name: str 

128 

129 with tempfile.NamedTemporaryFile(suffix=".sdf", delete_on_close=False) as tmp: 

130 tmp.close() 

131 with h5py.File(tmp.name, "w") as f: 

132 arch = NdfOutputArchive(f) 

133 ptr = arch.serialize_pointer("psf", lambda nested: TinyTree(name="hello"), key=("psf", 1)) 

134 with NdfInputArchive.open(tmp.name) as archive: 

135 # Deserializer just returns the model unchanged. 

136 result = archive.deserialize_pointer(ptr, TinyTree, lambda m, _a: m) 

137 self.assertEqual(result.name, "hello") 

138 

139 def test_deserialize_pointer_caches_by_ref(self): 

140 class TinyTree(pydantic.BaseModel): 

141 name: str 

142 

143 calls = [] 

144 

145 def deserializer(model, _archive): 

146 calls.append(model) 

147 return model 

148 

149 with tempfile.NamedTemporaryFile(suffix=".sdf", delete_on_close=False) as tmp: 

150 tmp.close() 

151 with h5py.File(tmp.name, "w") as f: 

152 arch = NdfOutputArchive(f) 

153 ptr = arch.serialize_pointer("psf", lambda nested: TinyTree(name="x"), key=("psf", 1)) 

154 with NdfInputArchive.open(tmp.name) as archive: 

155 first = archive.deserialize_pointer(ptr, TinyTree, deserializer) 

156 second = archive.deserialize_pointer(ptr, TinyTree, deserializer) 

157 self.assertIs(first, second) 

158 self.assertEqual(len(calls), 1) 

159 

160 def test_deserialize_pointer_caches_frame_set_for_get_frame_set(self): 

161 class TinyTree(pydantic.BaseModel): 

162 name: str 

163 

164 class DummyFrameSet(FrameSet): 

165 def __contains__(self, frame): 

166 return False 

167 

168 def __getitem__(self, key): 

169 raise AssertionError("DummyFrameSet should not be indexed in this test.") 

170 

171 sentinel = DummyFrameSet() 

172 

173 with tempfile.NamedTemporaryFile(suffix=".sdf", delete_on_close=False) as tmp: 

174 tmp.close() 

175 with h5py.File(tmp.name, "w") as f: 

176 arch = NdfOutputArchive(f) 

177 ptr = arch.serialize_frame_set( 

178 "frames", 

179 sentinel, 

180 lambda nested: TinyTree(name="frames"), 

181 key=("frames", 1), 

182 ) 

183 with NdfInputArchive.open(tmp.name) as archive: 

184 result = archive.deserialize_pointer(ptr, TinyTree, lambda _m, _a: sentinel) 

185 self.assertIs(result, sentinel) 

186 self.assertIs(archive.get_frame_set(ptr), sentinel) 

187 

188 def test_get_frame_set_returns_cached_value(self): 

189 # Exercise the cache mechanism with a sentinel object pretending 

190 # to be a FrameSet. Real FrameSet plumbing comes when the AST 

191 # text dump for /WCS/DATA lands in a follow-up task. 

192 sentinel = object() 

193 with tempfile.NamedTemporaryFile(suffix=".sdf", delete_on_close=False) as tmp: 

194 tmp.close() 

195 write(Image(np.zeros((2, 2), dtype=np.float32)), tmp.name) 

196 with NdfInputArchive.open(tmp.name) as archive: 

197 # Manually populate the cache as deserialize_pointer would 

198 # if a FrameSet deserializer ran. 

199 archive._frame_set_cache["/MORE/LSST/PIXEL_TO_SKY"] = sentinel 

200 pointer = NdfPointerModel(path="/MORE/LSST/PIXEL_TO_SKY") 

201 self.assertIs(archive.get_frame_set(pointer), sentinel) 

202 

203 def test_get_frame_set_raises_if_not_cached(self): 

204 with tempfile.NamedTemporaryFile(suffix=".sdf", delete_on_close=False) as tmp: 

205 tmp.close() 

206 write(Image(np.zeros((2, 2), dtype=np.float32)), tmp.name) 

207 with NdfInputArchive.open(tmp.name) as archive: 

208 pointer = NdfPointerModel(path="/MORE/LSST/UNKNOWN") 

209 with self.assertRaises(AssertionError): 

210 archive.get_frame_set(pointer) 

211 

212 

213@unittest.skipUnless(HAVE_H5PY, "h5py is not installed") 

214class NdfInputArchiveOpaqueMetadataTestCase(unittest.TestCase): 

215 """Tests for `NdfInputArchive.get_opaque_metadata`.""" 

216 

217 def test_more_fits_round_trips_via_opaque_metadata(self): 

218 image = Image(np.zeros((2, 2), dtype=np.float32)) 

219 primary = astropy.io.fits.Header() 

220 primary["FOO"] = ("bar", "test card") 

221 opaque = FitsOpaqueMetadata() 

222 opaque.add_header(primary, name="", ver=1) 

223 image._opaque_metadata = opaque 

224 with tempfile.NamedTemporaryFile(suffix=".sdf", delete_on_close=False) as tmp: 

225 tmp.close() 

226 write(image, tmp.name) 

227 with NdfInputArchive.open(tmp.name) as archive: 

228 recovered = archive.get_opaque_metadata() 

229 self.assertIn(ExtensionKey(), recovered.headers) 

230 self.assertEqual(recovered.headers[ExtensionKey()]["FOO"], "bar") 

231 

232 def test_get_opaque_metadata_empty_when_no_more_fits(self): 

233 # Image with no opaque metadata -> /MORE/FITS is absent in the file. 

234 image = Image(np.zeros((2, 2), dtype=np.float32)) 

235 with tempfile.NamedTemporaryFile(suffix=".sdf", delete_on_close=False) as tmp: 

236 tmp.close() 

237 write(image, tmp.name) 

238 with NdfInputArchive.open(tmp.name) as archive: 

239 recovered = archive.get_opaque_metadata() 

240 self.assertIsInstance(recovered, FitsOpaqueMetadata) 

241 # No primary header should be populated since /MORE/FITS 

242 # was never written. 

243 self.assertFalse(recovered.headers) 

244 

245 

246@unittest.skipUnless(HAVE_H5PY, "h5py is not installed") 

247class NdfReadFunctionTestCase(unittest.TestCase): 

248 """Tests for the generic ``read`` (symmetric LSST trees) and the 

249 NDF-specific ``ndf.read_starlink`` (schema-less auto-detect). 

250 """ 

251 

252 def test_read_round_trips_image(self): 

253 image = Image( 

254 np.arange(20, dtype=np.float32).reshape(4, 5), 

255 bbox=Box.factory[10:14, 20:25], 

256 ) 

257 with tempfile.NamedTemporaryFile(suffix=".sdf", delete_on_close=False) as tmp: 

258 tmp.close() 

259 write(image, tmp.name) 

260 result = read(tmp.name, Image) 

261 self.assertIsInstance(result, Image) 

262 np.testing.assert_array_equal(result.array, image.array) 

263 self.assertEqual(result.bbox, image.bbox) 

264 

265 def test_read_starlink_file_auto_detects_image(self): 

266 # The canonical fixture has no /MORE/LSST/JSON, no QUALITY, 

267 # no VARIANCE -- auto-detect should return an Image whose array 

268 # shape matches the file (611x609 int16). 

269 example_path = os.path.join(os.path.dirname(__file__), "data", "example-ndf.sdf") 

270 result = read_starlink(Image, example_path) 

271 self.assertIsInstance(result, Image) 

272 self.assertEqual(result.array.shape, (611, 609)) 

273 self.assertEqual(result.array.dtype, np.int16) 

274 self.assertIsNotNone(result.sky_projection) 

275 

276 def test_read_starlink_file_recovers_opaque_fits_metadata(self): 

277 example_path = os.path.join(os.path.dirname(__file__), "data", "example-ndf.sdf") 

278 result = read_starlink(Image, example_path) 

279 opaque = result._opaque_metadata 

280 self.assertIn(ExtensionKey(), opaque.headers) 

281 # The fixture is a real Starlink M57 image; sample one card we know 

282 # is present (NAXIS). 

283 primary = opaque.headers[ExtensionKey()] 

284 self.assertIn("NAXIS", primary) 

285 

286 def test_read_auto_detects_nested_quality_array(self): 

287 image_array = np.arange(6, dtype=np.float32).reshape(2, 3) 

288 quality_array = np.array([[0, 1, 0], [1, 0, 1]], dtype=np.uint8) 

289 

290 with tempfile.NamedTemporaryFile(suffix=".sdf", delete_on_close=False) as tmp: 

291 tmp.close() 

292 with h5py.File(tmp.name, "w") as f: 

293 _hds.set_root_name(f, "TEST", "NDF") 

294 data_array = _hds.create_structure(f, "DATA_ARRAY", "ARRAY") 

295 _hds.write_array(data_array, "DATA", image_array) 

296 quality = _hds.create_structure(f, "QUALITY", "QUALITY") 

297 quality_array_struct = _hds.create_structure(quality, "QUALITY", "ARRAY") 

298 _hds.write_array(quality_array_struct, "DATA", quality_array) 

299 _hds.write_array(quality_array_struct, "ORIGIN", np.array([0, 0], dtype=np.int32)) 

300 _hds.write_array(quality_array_struct, "BAD_PIXEL", np.array(False, dtype=np.bool_)) 

301 _hds.write_array(quality, "BADBITS", np.array(1, dtype=np.uint8)) 

302 result = read_starlink(MaskedImage, tmp.name) 

303 self.assertIsInstance(result, MaskedImage) 

304 np.testing.assert_array_equal(result.mask.array[:, :, 0], quality_array) 

305 self.assertEqual(set(result.mask.schema.names), {f"MASK{i}" for i in range(8)}) 

306 image_result = read_starlink(Image, tmp.name) 

307 self.assertIsInstance(image_result, Image) 

308 np.testing.assert_array_equal(image_result.array, image_array) 

309 

310 def test_read_auto_detect_preserves_quality_bits(self): 

311 image_array = np.arange(6, dtype=np.float32).reshape(2, 3) 

312 quality_array = np.array([[0, 2, 4], [2, 0, 6]], dtype=np.uint8) 

313 expected_mask1 = np.array([[0, 1, 0], [1, 0, 1]], dtype=bool) 

314 expected_mask2 = np.array([[0, 0, 1], [0, 0, 1]], dtype=bool) 

315 

316 with tempfile.NamedTemporaryFile(suffix=".sdf", delete_on_close=False) as tmp: 

317 tmp.close() 

318 with h5py.File(tmp.name, "w") as f: 

319 _hds.set_root_name(f, "TEST", "NDF") 

320 data_array = _hds.create_structure(f, "DATA_ARRAY", "ARRAY") 

321 _hds.write_array(data_array, "DATA", image_array) 

322 quality = _hds.create_structure(f, "QUALITY", "QUALITY") 

323 quality_array_struct = _hds.create_structure(quality, "QUALITY", "ARRAY") 

324 _hds.write_array(quality_array_struct, "DATA", quality_array) 

325 _hds.write_array(quality_array_struct, "ORIGIN", np.array([0, 0], dtype=np.int32)) 

326 _hds.write_array(quality_array_struct, "BAD_PIXEL", np.array(False, dtype=np.bool_)) 

327 _hds.write_array(quality, "BADBITS", np.array(2, dtype=np.uint8)) 

328 result = read_starlink(MaskedImage, tmp.name) 

329 self.assertIsInstance(result, MaskedImage) 

330 mask = result.mask 

331 np.testing.assert_array_equal(mask.array[:, :, 0], quality_array) 

332 np.testing.assert_array_equal(mask.get("MASK1"), expected_mask1) 

333 np.testing.assert_array_equal(mask.get("MASK2"), expected_mask2) 

334 self.assertIn("Selected by BADBITS", mask.schema.descriptions["MASK1"]) 

335 self.assertNotIn("Selected by BADBITS", mask.schema.descriptions["MASK2"]) 

336 

337 def test_read_auto_detected_data_only_as_masked_image_uses_defaults(self): 

338 image_array = np.arange(6, dtype=np.float32).reshape(2, 3) 

339 

340 with tempfile.NamedTemporaryFile(suffix=".sdf", delete_on_close=False) as tmp: 

341 tmp.close() 

342 with h5py.File(tmp.name, "w") as f: 

343 _hds.set_root_name(f, "TEST", "NDF") 

344 data_array = _hds.create_structure(f, "DATA_ARRAY", "ARRAY") 

345 _hds.write_array(data_array, "DATA", image_array) 

346 _hds.write_array(data_array, "ORIGIN", np.array([5, 4], dtype=np.int32)) 

347 result = read_starlink(MaskedImage, tmp.name) 

348 self.assertIsInstance(result, MaskedImage) 

349 self.assertEqual(result.bbox, Box.factory[4:6, 5:8]) 

350 np.testing.assert_array_equal(result.image.array, image_array) 

351 np.testing.assert_array_equal( 

352 result.mask.array, 

353 np.zeros((2, 3, 1), dtype=np.uint8), 

354 ) 

355 np.testing.assert_array_equal( 

356 result.variance.array, 

357 np.ones((2, 3), dtype=np.float32), 

358 ) 

359 

360 def test_read_auto_detected_variance_as_masked_image_keeps_variance(self): 

361 image_array = np.arange(6, dtype=np.float32).reshape(2, 3) 

362 variance_array = np.full((2, 3), 2.5, dtype=np.float32) 

363 

364 with tempfile.NamedTemporaryFile(suffix=".sdf", delete_on_close=False) as tmp: 

365 tmp.close() 

366 with h5py.File(tmp.name, "w") as f: 

367 _hds.set_root_name(f, "TEST", "NDF") 

368 data_array = _hds.create_structure(f, "DATA_ARRAY", "ARRAY") 

369 _hds.write_array(data_array, "DATA", image_array) 

370 _hds.write_array(data_array, "ORIGIN", np.array([5, 4], dtype=np.int32)) 

371 variance = _hds.create_structure(f, "VARIANCE", "ARRAY") 

372 _hds.write_array(variance, "DATA", variance_array) 

373 _hds.write_array(variance, "ORIGIN", np.array([5, 4], dtype=np.int32)) 

374 result = read_starlink(MaskedImage, tmp.name) 

375 self.assertIsInstance(result, MaskedImage) 

376 np.testing.assert_array_equal(result.variance.array, variance_array) 

377 np.testing.assert_array_equal( 

378 result.mask.array, 

379 np.zeros((2, 3, 1), dtype=np.uint8), 

380 ) 

381 

382 def test_read_auto_detected_units_component(self): 

383 image_array = np.arange(6, dtype=np.float32).reshape(2, 3) 

384 

385 with tempfile.NamedTemporaryFile(suffix=".sdf", delete_on_close=False) as tmp: 

386 tmp.close() 

387 with h5py.File(tmp.name, "w") as f: 

388 _hds.set_root_name(f, "TEST", "NDF") 

389 data_array = _hds.create_structure(f, "DATA_ARRAY", "ARRAY") 

390 _hds.write_array(data_array, "DATA", image_array) 

391 f.create_dataset("UNITS", data=np.bytes_("count")) 

392 result = read_starlink(Image, tmp.name) 

393 self.assertEqual(result.unit, u.ct) 

394 

395 def test_read_missing_data_array_raises(self): 

396 # A file with only /MORE/LSST/JSON is fine for the symmetric 

397 # path. A file with NEITHER /MORE/LSST/JSON NOR DATA_ARRAY is a 

398 # malformed NDF -- auto-detect must fail clearly. 

399 with tempfile.NamedTemporaryFile(suffix=".sdf", delete_on_close=False) as tmp: 

400 tmp.close() 

401 with h5py.File(tmp.name, "w") as f: 

402 f["/"].attrs["CLASS"] = "NDF" 

403 # Note: no DATA_ARRAY, no /MORE/LSST/JSON. 

404 with self.assertRaises(ArchiveReadError): 

405 read_starlink(Image, tmp.name) 

406 

407 def test_read_auto_detect_wrong_target_type_raises(self): 

408 # Auto-detect only knows how to produce Image-like objects from NDF 

409 # components; unrelated target classes should fail clearly. 

410 example_path = os.path.join(os.path.dirname(__file__), "data", "example-ndf.sdf") 

411 with self.assertRaises(ArchiveReadError): 

412 read_starlink(Mask, example_path)