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.
12from __future__ import annotations
14import os
15import tempfile
16import unittest
18import astropy.io.fits
19import astropy.units as u
20import numpy as np
21import pydantic
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)
34try:
35 import h5py
37 from lsst.images.ndf import (
38 NdfInputArchive,
39 NdfOutputArchive,
40 NdfPointerModel,
41 _hds,
42 read_starlink,
43 write,
44 )
46 HAVE_H5PY = True
47except ImportError:
48 HAVE_H5PY = False
51@unittest.skipUnless(HAVE_H5PY, "h5py is not installed")
52class NdfInputArchiveOpenTestCase(unittest.TestCase):
53 """Tests for `NdfInputArchive.open` and `get_tree`."""
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())
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)
79@unittest.skipUnless(HAVE_H5PY, "h5py is not installed")
80class NdfInputArchiveDataTestCase(unittest.TestCase):
81 """Tests for `get_array`, `deserialize_pointer`, and `get_frame_set`."""
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)
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])
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]))
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)
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
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")
139 def test_deserialize_pointer_caches_by_ref(self):
140 class TinyTree(pydantic.BaseModel):
141 name: str
143 calls = []
145 def deserializer(model, _archive):
146 calls.append(model)
147 return model
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)
160 def test_deserialize_pointer_caches_frame_set_for_get_frame_set(self):
161 class TinyTree(pydantic.BaseModel):
162 name: str
164 class DummyFrameSet(FrameSet):
165 def __contains__(self, frame):
166 return False
168 def __getitem__(self, key):
169 raise AssertionError("DummyFrameSet should not be indexed in this test.")
171 sentinel = DummyFrameSet()
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)
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)
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)
213@unittest.skipUnless(HAVE_H5PY, "h5py is not installed")
214class NdfInputArchiveOpaqueMetadataTestCase(unittest.TestCase):
215 """Tests for `NdfInputArchive.get_opaque_metadata`."""
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")
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)
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 """
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)
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)
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)
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)
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)
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)
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"])
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)
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 )
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)
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 )
382 def test_read_auto_detected_units_component(self):
383 image_array = np.arange(6, dtype=np.float32).reshape(2, 3)
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)
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)
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)