Coverage for tests/test_mask.py: 93%
307 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-16 15:24 -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 dataclasses
15import os
16from typing import Any
18import astropy.io.fits
19import numpy as np
20import pytest
22import lsst.utils.tests
23from lsst.images import (
24 Box,
25 Mask,
26 MaskPlane,
27 MaskSchema,
28 get_legacy_non_cell_coadd_mask_planes,
29 get_legacy_visit_image_mask_planes,
30)
31from lsst.images._mask import _guess_legacy_plane_map
32from lsst.images.tests import RoundtripFits, assert_masks_equal, compare_mask_to_legacy
34try:
35 from lsst.afw.image import MaskedImageReader as LegacyMaskedImageReader
37except ImportError:
38 type LegacyMaskedImageReader = Any # type: ignore[no-redef]
40EXTERNAL_DATA_DIR = os.environ.get("TESTDATA_IMAGES_DIR", None)
43@dataclasses.dataclass
44class _LegacyTestData:
45 mask: Mask
46 reader: LegacyMaskedImageReader
47 plane_map: dict[str, MaskPlane]
50@pytest.fixture(scope="session")
51def legacy_test_data() -> _LegacyTestData:
52 """Return a Mask read directly from the legacy test dataset and a legacy
53 reader for that image.
55 Skips if TESTDATA_IMAGES_DIR is unset or lsst.afw.image is unavailable.
56 """
57 if EXTERNAL_DATA_DIR is None: 57 ↛ 59line 57 didn't jump to line 59 because the condition on line 57 was always true
58 pytest.skip("TESTDATA_IMAGES_DIR is not in the environment.")
59 try:
60 from lsst.afw.image import MaskedImageFitsReader
61 except ImportError:
62 pytest.skip("'lsst.afw.image' could not be imported.")
63 filename = os.path.join(EXTERNAL_DATA_DIR, "dp2", "legacy", "visit_image.fits")
64 plane_map = get_legacy_visit_image_mask_planes()
65 mask = Mask.read_legacy(filename, ext=2, plane_map=plane_map)
66 reader = MaskedImageFitsReader(filename)
67 return _LegacyTestData(mask=mask, reader=reader, plane_map=plane_map)
70def make_mask_planes(rng: np.random.Generator, n_planes: int, n_placeholders: int) -> list[MaskPlane | None]:
71 """Return a shuffled list of MaskPlane objects with placeholder Nones."""
72 planes: list[MaskPlane | None] = []
73 for i in range(n_planes):
74 planes.append(MaskPlane(f"M{i}", f"D{i}"))
75 planes.extend([None] * n_placeholders)
76 rng.shuffle(planes)
77 return planes
80def test_schema() -> None:
81 """Test MaskSchema construction, accessors, and basic operations."""
82 rng = np.random.default_rng(500)
83 planes = make_mask_planes(rng, 17, 5)
84 with pytest.raises(TypeError):
85 MaskSchema.bits_per_element(np.float32)
86 assert MaskSchema.bits_per_element(np.uint8) == 8
87 schema = MaskSchema(planes, dtype=np.uint8)
88 assert list(schema) == planes
89 assert len(schema) == len(planes)
90 assert schema[5] == planes[5]
91 assert eval(repr(schema), {"dtype": np.dtype, "MaskSchema": MaskSchema, "MaskPlane": MaskPlane}) == schema
92 string = str(schema)
93 assert len(string.split("\n")) == 17
94 bit5 = schema.bit("M5")
95 assert f"M5 [{bit5.index}@{hex(bit5.mask)}]: D5" in string
96 assert schema == MaskSchema(planes, np.uint8)
97 assert schema != MaskSchema(planes, np.int16)
98 assert schema != MaskSchema(planes[:-1], np.uint8)
99 assert schema.dtype == np.dtype(np.uint8)
100 assert schema.mask_size == 3
101 assert schema.names == {f"M{i}" for i in range(17)}
102 assert schema.descriptions == {f"M{i}": f"D{i}" for i in range(17)}
103 bit7 = schema.bit("M7")
104 bitmask57 = schema.bitmask("M5", "M7")
105 assert bitmask57[bit5.index] & bit5.mask
106 assert bitmask57[bit7.index] & bit7.mask
107 bitmask57[bit5.index] &= ~bit5.mask
108 bitmask57[bit7.index] &= ~bit7.mask
109 assert not bitmask57.any()
110 splits = schema.split(np.int16)
111 assert len(splits) == 2
112 assert splits[0].mask_size == 1
113 assert splits[1].mask_size == 1
114 assert list(splits[0]) + list(splits[1]) == [p for p in planes if p is not None]
115 assert len(splits[0]) == 15
116 assert len(splits[1]) == 2
119def test_schema_from_fits_header() -> None:
120 """Verify MaskSchema.from_fits_header inverts update_header."""
121 planes = [
122 MaskPlane("NO_DATA", "No data was available for this pixel."),
123 MaskPlane("COSMIC_RAY", "A cosmic ray affected this pixel."),
124 MaskPlane("DETECTED", "Pixel was part of a detected source."),
125 ]
126 schema = MaskSchema(planes, dtype=np.uint8)
127 header = astropy.io.fits.Header()
128 schema.update_header(header)
129 result = MaskSchema.from_fits_header(header)
130 assert result.dtype == np.dtype(np.uint8)
131 assert list(result) == planes
132 assert result == schema
135def test_schema_from_fits_header_preserves_gaps() -> None:
136 """Verify None placeholders are reconstructed from gaps in MSKN card
137 numbering.
138 """
139 planes: list[MaskPlane | None] = [MaskPlane("A", "a"), None, MaskPlane("B", "b")]
140 header = astropy.io.fits.Header()
141 MaskSchema(planes, dtype=np.uint8).update_header(header)
142 assert list(MaskSchema.from_fits_header(header)) == planes
145def test_schema_from_fits_header_requires_cards() -> None:
146 """Verify MaskSchema.from_fits_header raises ValueError on a header with
147 no MSKN cards.
148 """
149 with pytest.raises(ValueError):
150 MaskSchema.from_fits_header(astropy.io.fits.Header())
153def test_interpret() -> None:
154 """Verify interpret returns correct plane names across a multi-byte
155 schema.
156 """
157 # 3 named planes padded with Nones to exceed 8 bits; A and B land in
158 # byte 0, C lands in byte 1.
159 planes: list[MaskPlane | None] = [
160 MaskPlane("A", "a"),
161 MaskPlane("B", "b"),
162 None,
163 None,
164 None,
165 None,
166 None,
167 None,
168 MaskPlane("C", "c"),
169 ]
170 schema = MaskSchema(planes, dtype=np.uint8)
171 assert schema.mask_size == 2
173 assert set(schema.interpret(schema.bitmask("A", "C"))) == {"A", "C"}
174 assert set(schema.interpret(schema.bitmask("B"))) == {"B"}
175 assert set(schema.interpret(np.zeros(schema.mask_size, dtype=schema.dtype))) == set()
176 assert set(schema.interpret(schema.bitmask("A", "B", "C"))) == {"A", "B", "C"}
179def test_basics() -> None:
180 """Test basic Mask construction, string representation, and error
181 conditions.
182 """
183 rng = np.random.default_rng(500)
184 planes = make_mask_planes(rng, 35, n_placeholders=5)
185 schema = MaskSchema(planes, dtype=np.uint8)
186 bbox = Box.factory[5:50, 6:60]
187 mask = Mask(
188 0,
189 schema=schema,
190 bbox=bbox,
191 metadata={"four_and_a_half": 4.5},
192 )
194 assert mask[...] is not mask
195 assert mask.__eq__(42) == NotImplemented
196 assert mask == mask
197 assert (
198 str(mask)
199 == "Mask([y=5:50, x=6:60], ['M34', 'M15', 'M29', 'M1', 'M20', 'M11', 'M13', 'M7', 'M17', 'M12', "
200 "'M31', 'M16', 'M2', 'M3', 'M8', 'M26', 'M22', 'M5', 'M18', 'M19', 'M24', 'M21', 'M27', 'M6', "
201 "'M28', 'M10', 'M4', 'M23', 'M0', 'M25', 'M9', 'M14', 'M33', 'M32', 'M30'])"
202 )
203 assert repr(mask).startswith(
204 "Mask(..., bbox=Box(y=Interval(start=5, stop=50), x=Interval(start=6, stop=60)), "
205 "schema=MaskSchema([MaskPlane(name='M34', description='D34')"
206 ), f"Repr: {mask!r}"
208 with pytest.raises(TypeError):
209 # No bbox, size or array.
210 Mask(0, schema=schema)
212 with pytest.raises(ValueError):
213 # Box mismatch.
214 Mask(mask.array, schema=schema, bbox=Box.factory[0:20, -5:45])
216 with pytest.raises(ValueError):
217 # Shape mismatch.
218 Mask(mask.array, schema=schema, shape=(5, 10, 5))
220 with pytest.raises(ValueError):
221 # Cannot be 2-D.
222 Mask(mask.array.reshape((2430, 5)), schema=schema, bbox=Box.factory[0:20, -5:45])
225def test_read_write() -> None:
226 """Test explicit calls to Mask.read and Mask.write through FITS."""
227 rng = np.random.default_rng(500)
228 planes = make_mask_planes(rng, 35, n_placeholders=5)
229 schema = MaskSchema(planes, dtype=np.uint8)
230 bbox = Box.factory[5:50, 6:60]
231 mask = Mask(
232 0,
233 schema=schema,
234 bbox=bbox,
235 metadata={"four_and_a_half": 4.5},
236 )
237 with lsst.utils.tests.getTempFilePath(".fits") as tmpFile:
238 mask.write(tmpFile)
239 new = Mask.read(tmpFile)
240 assert new == mask
241 # __eq__ ignores metadata.
242 assert new.metadata["four_and_a_half"] == 4.5
243 assert new.metadata == mask.metadata
246def test_serialize_multi() -> None:
247 """Test serializing a mask with more than 31 mask planes (multiple
248 HDUs).
249 """
250 rng = np.random.default_rng(500)
251 planes = make_mask_planes(rng, 35, n_placeholders=5)
252 schema = MaskSchema(planes, dtype=np.uint8)
253 bbox = Box.factory[5:50, 6:60]
254 mask = Mask(0, schema=schema, bbox=bbox, metadata={"four_and_a_half": 4.5})
255 shape = bbox.shape
256 for plane in schema:
257 if plane is not None:
258 mask.set(plane.name, rng.random(shape) > 0.5)
259 with RoundtripFits(mask) as roundtrip:
260 fits = roundtrip.inspect()
261 assert fits[1].header["EXTNAME"] == "MASK"
262 assert fits[1].header.get("EXTVER", 1) == 1
263 assert fits[1].header["ZCMPTYPE"] == "GZIP_2"
264 assert fits[2].header["EXTNAME"] == "MASK"
265 assert fits[2].header["EXTVER"] == 2
266 assert fits[2].header["ZCMPTYPE"] == "GZIP_2"
267 n = 0
268 for plane in planes:
269 if plane is not None:
270 hdu = fits[1] if n < 31 else fits[2]
271 assert hdu.header[f"MSKN{(n % 31):04d}"] == plane.name
272 assert hdu.header[f"MSKM{(n % 31):04d}"] == 1 << (n % 31)
273 assert hdu.header[f"MSKD{(n % 31):04d}"] == plane.description
274 n += 1
275 assert_masks_equal(mask, roundtrip.result)
278def test_add_plane_returns_new_mask() -> None:
279 """Verify add_plane returns a new mask without modifying the original or
280 its views.
281 """
282 rng = np.random.default_rng(500)
283 planes = make_mask_planes(rng, 3, n_placeholders=0)
284 schema = MaskSchema(planes, dtype=np.uint8)
285 bbox = Box.factory[5:50, 6:60]
286 mask = Mask(0, schema=schema, bbox=bbox)
287 m0 = rng.random(bbox.shape) > 0.5
288 mask.set("M0", m0)
289 view = mask[bbox] # shares the array and old schema with mask
290 original_array = mask.array
292 new_mask = mask.add_plane("OUTSIDE_STENCIL", "Pixel lies outside the stencil.")
294 # The original mask and any views keep the old schema and array.
295 assert "OUTSIDE_STENCIL" not in mask.schema.names
296 assert "OUTSIDE_STENCIL" not in view.schema.names
297 assert mask.array is original_array
298 # The new mask reallocated a fresh array and carries the new plane.
299 assert new_mask.array is not original_array
300 assert "OUTSIDE_STENCIL" in new_mask.schema.names
301 assert new_mask.schema.descriptions["OUTSIDE_STENCIL"] == "Pixel lies outside the stencil."
302 # The new plane is the fourth (overall index 3) so it lives in byte 0.
303 bit = new_mask.schema.bit("OUTSIDE_STENCIL")
304 assert bit.index == 0
305 assert bit.mask == 1 << 3
306 assert new_mask.schema.mask_size == 1
307 # Existing plane data is preserved and the new plane starts all-False.
308 np.testing.assert_array_equal(new_mask.get("M0"), m0)
309 assert not new_mask.get("OUTSIDE_STENCIL").any()
312def test_add_plane_grows_byte() -> None:
313 """Verify adding a ninth plane crosses the 8-plane boundary into a
314 second byte.
315 """
316 rng = np.random.default_rng(500)
317 planes = make_mask_planes(rng, 8, n_placeholders=0)
318 schema = MaskSchema(planes, dtype=np.uint8)
319 bbox = Box.factory[5:50, 6:60]
320 mask = Mask(0, schema=schema, bbox=bbox)
321 set_planes = {}
322 for plane in planes:
323 assert plane is not None
324 boolean_mask = rng.random(bbox.shape) > 0.5
325 mask.set(plane.name, boolean_mask)
326 set_planes[plane.name] = boolean_mask
328 new_mask = mask.add_plane("OUTSIDE_STENCIL", "Pixel lies outside the stencil.")
330 # The original is unchanged; the new mask spills into a second byte.
331 assert mask.schema.mask_size == 1
332 bit = new_mask.schema.bit("OUTSIDE_STENCIL")
333 assert bit.index == 1
334 assert bit.mask == 1 << 0
335 assert new_mask.schema.mask_size == 2
336 assert new_mask.array.shape == bbox.shape + (2,)
337 assert not new_mask.get("OUTSIDE_STENCIL").any()
338 # Every pre-existing plane keeps its data.
339 for name, boolean_mask in set_planes.items():
340 np.testing.assert_array_equal(new_mask.get(name), boolean_mask)
343def test_add_planes_multiple() -> None:
344 """Verify add_planes adds several planes in a single call."""
345 rng = np.random.default_rng(500)
346 planes = make_mask_planes(rng, 3, n_placeholders=0)
347 bbox = Box.factory[0:4, 0:5]
348 mask = Mask(0, schema=MaskSchema(planes, dtype=np.uint8), bbox=bbox)
349 m0 = rng.random(bbox.shape) > 0.5
350 mask.set("M0", m0)
352 new_mask = mask.add_planes([MaskPlane("A", "plane a"), MaskPlane("B", "plane b")])
354 assert set(mask.schema.names) == {"M0", "M1", "M2"} # original unchanged
355 assert set(new_mask.schema.names) == {"M0", "M1", "M2", "A", "B"}
356 np.testing.assert_array_equal(new_mask.get("M0"), m0)
357 assert not new_mask.get("A").any()
358 assert not new_mask.get("B").any()
361def test_add_planes_drop_reassigns_bits() -> None:
362 """Verify dropping a plane compacts the schema and repacks pixel values."""
363 rng = np.random.default_rng(500)
364 bbox = Box.factory[0:4, 0:5]
365 schema = MaskSchema([MaskPlane("A", "a"), MaskPlane("B", "b"), MaskPlane("C", "c")], dtype=np.uint8)
366 mask = Mask(0, schema=schema, bbox=bbox)
367 a = rng.random(bbox.shape) > 0.5
368 c = rng.random(bbox.shape) > 0.5
369 mask.set("A", a)
370 mask.set("B", rng.random(bbox.shape) > 0.5)
371 mask.set("C", c)
373 new_mask = mask.add_planes([MaskPlane("D", "d")], drop=["B"])
375 # B is gone; D is appended after the retained planes.
376 assert list(new_mask.schema.names) == ["A", "C", "D"]
377 assert "B" not in new_mask.schema.names
378 # C moved down from bit 2 to bit 1; D takes bit 2.
379 assert new_mask.schema.bit("A").mask == 1 << 0
380 assert new_mask.schema.bit("C").mask == 1 << 1
381 assert new_mask.schema.bit("D").mask == 1 << 2
382 # Retained pixel values follow their planes; the new plane is cleared.
383 np.testing.assert_array_equal(new_mask.get("A"), a)
384 np.testing.assert_array_equal(new_mask.get("C"), c)
385 assert not new_mask.get("D").any()
388def test_add_planes_with_placeholder() -> None:
389 """Verify None placeholders reserve bits and survive add_planes and a
390 FITS round-trip.
391 """
392 rng = np.random.default_rng(500)
393 bbox = Box.factory[0:4, 0:5]
394 # Schema with a pre-existing placeholder reserving bit 1.
395 schema = MaskSchema([MaskPlane("A", "a"), None, MaskPlane("B", "b")], dtype=np.uint8)
396 mask = Mask(0, schema=schema, bbox=bbox)
397 a = rng.random(bbox.shape) > 0.5
398 b = rng.random(bbox.shape) > 0.5
399 mask.set("A", a)
400 mask.set("B", b)
402 # Append a block that itself contains an interior placeholder.
403 new_mask = mask.add_planes([MaskPlane("C", "c"), None, MaskPlane("D", "d")])
405 # The pre-existing placeholder stays at bit 1; the added placeholder
406 # stays between C and D (bit 4), not at the end.
407 assert list(new_mask.schema) == [
408 MaskPlane("A", "a"),
409 None,
410 MaskPlane("B", "b"),
411 MaskPlane("C", "c"),
412 None,
413 MaskPlane("D", "d"),
414 ]
415 assert new_mask.schema.bit("A").mask == 1 << 0
416 assert new_mask.schema.bit("B").mask == 1 << 2
417 assert new_mask.schema.bit("C").mask == 1 << 3
418 assert new_mask.schema.bit("D").mask == 1 << 5
419 # Retained pixel values follow their planes; new planes start cleared.
420 np.testing.assert_array_equal(new_mask.get("A"), a)
421 np.testing.assert_array_equal(new_mask.get("B"), b)
422 assert not new_mask.get("C").any()
423 assert not new_mask.get("D").any()
425 with RoundtripFits(new_mask) as roundtrip:
426 assert_masks_equal(new_mask, roundtrip.result)
429def test_add_planes_drop_unknown_raises() -> None:
430 """Verify dropping a non-existent plane raises ValueError."""
431 mask = Mask(0, schema=MaskSchema([MaskPlane("A", "a")], dtype=np.uint8), bbox=Box.factory[0:2, 0:2])
432 with pytest.raises(ValueError):
433 mask.add_planes([], drop=["NOPE"])
436def test_add_plane_duplicate_raises() -> None:
437 """Verify adding a plane whose name already exists raises ValueError."""
438 rng = np.random.default_rng(500)
439 planes = make_mask_planes(rng, 3, n_placeholders=0)
440 schema = MaskSchema(planes, dtype=np.uint8)
441 mask = Mask(0, schema=schema, bbox=Box.factory[0:4, 0:4])
442 with pytest.raises(ValueError):
443 mask.add_plane("M0", "Duplicate of an existing plane.")
446def test_add_plane_roundtrip() -> None:
447 """Verify a runtime-added plane and its data survive a FITS round-trip."""
448 rng = np.random.default_rng(500)
449 planes = make_mask_planes(rng, 8, n_placeholders=0)
450 schema = MaskSchema(planes, dtype=np.uint8)
451 bbox = Box.factory[5:50, 6:60]
452 mask = Mask(0, schema=schema, bbox=bbox)
453 mask = mask.add_plane("OUTSIDE_STENCIL", "Pixel lies outside the stencil.")
454 mask.set("OUTSIDE_STENCIL", rng.random(bbox.shape) > 0.5)
455 with lsst.utils.tests.getTempFilePath(".fits") as tmpFile:
456 mask.write(tmpFile)
457 new = Mask.read(tmpFile)
458 assert new == mask
459 assert new.schema.descriptions["OUTSIDE_STENCIL"] == "Pixel lies outside the stencil."
460 assert_masks_equal(new, mask)
463def test_legacy_non_cell_coadd_plane_map() -> None:
464 """Verify the non-cell coadd map defines a distinct SENSOR_EDGE plane."""
465 plane_map = get_legacy_non_cell_coadd_mask_planes()
466 assert "SENSOR_EDGE" in plane_map
467 assert plane_map["SENSOR_EDGE"].name == "SENSOR_EDGE"
470def test_guess_legacy_plane_map_coadd_discriminator() -> None:
471 """Verify INEXACT_PSF routes to a coadd map and SENSOR_EDGE discriminates
472 non-cell from cell.
473 """
474 non_cell = _guess_legacy_plane_map({"INEXACT_PSF": 11, "SENSOR_EDGE": 14})
475 assert "SENSOR_EDGE" in non_cell
476 cell = _guess_legacy_plane_map({"INEXACT_PSF": 11})
477 assert "SENSOR_EDGE" not in cell
480def test_legacy(legacy_test_data: _LegacyTestData) -> None:
481 """Test Mask.read_legacy, Mask.to_legacy, and Mask.from_legacy."""
482 assert legacy_test_data.mask.schema.names == {p.name for p in legacy_test_data.plane_map.values()}
483 assert legacy_test_data.mask.bbox == Box.from_legacy(legacy_test_data.reader.readBBox())
484 legacy_mask = legacy_test_data.reader.readMask()
485 compare_mask_to_legacy(legacy_test_data.mask, legacy_mask, legacy_test_data.plane_map)
486 compare_mask_to_legacy(
487 legacy_test_data.mask,
488 legacy_test_data.mask.to_legacy(legacy_test_data.plane_map),
489 legacy_test_data.plane_map,
490 )
491 assert_masks_equal(
492 legacy_test_data.mask, Mask.from_legacy(legacy_mask, plane_map=legacy_test_data.plane_map)
493 )
494 # Write the mask out in the new format, and test that we can read it back.
495 with RoundtripFits(legacy_test_data.mask, storage_class="MaskV2") as roundtrip:
496 pass
497 assert_masks_equal(roundtrip.result, legacy_test_data.mask)
500def test_legacy_butler_read(legacy_test_data: _LegacyTestData) -> None:
501 """Test that a round-tripped MaskV2 can be read back as a legacy afw
502 Mask via Butler.
503 """
504 with RoundtripFits(legacy_test_data.mask, storage_class="MaskV2") as roundtrip:
505 legacy_mask = roundtrip.get(storageClass="Mask")
506 assert isinstance(legacy_mask, lsst.afw.image.Mask)
507 compare_mask_to_legacy(legacy_test_data.mask, legacy_mask)