Coverage for tests/test_mask.py: 92%
281 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-08 02:03 -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 unittest
17import astropy.io.fits
18import numpy as np
20import lsst.utils.tests
21from lsst.images import (
22 Box,
23 Mask,
24 MaskPlane,
25 MaskSchema,
26 get_legacy_non_cell_coadd_mask_planes,
27 get_legacy_visit_image_mask_planes,
28)
29from lsst.images._mask import _guess_legacy_plane_map
30from lsst.images.tests import RoundtripFits, assert_masks_equal, compare_mask_to_legacy
32DATA_DIR = os.environ.get("TESTDATA_IMAGES_DIR", None)
35class MaskTestCase(unittest.TestCase):
36 """Tests for Mask and its helper classes."""
38 def setUp(self) -> None:
39 self.maxDiff = None
40 self.rng = np.random.default_rng(500)
42 def make_mask_planes(self, n_planes: int, n_placeholders: int) -> list[MaskPlane | None]:
43 planes: list[MaskPlane | None] = []
44 for i in range(n_planes):
45 planes.append(MaskPlane(f"M{i}", f"D{i}"))
46 planes.extend([None] * n_placeholders)
47 self.rng.shuffle(planes)
48 return planes
50 def test_schema(self) -> None:
51 """Test MaskSchema."""
52 self.assertEqual(MaskSchema.bits_per_element(np.uint8), 8)
53 planes = self.make_mask_planes(17, 5)
54 with self.assertRaises(TypeError):
55 MaskSchema.bits_per_element(np.float32)
56 schema = MaskSchema(planes, dtype=np.uint8)
57 self.assertEqual(list(schema), planes)
58 self.assertEqual(len(schema), len(planes))
59 self.assertEqual(schema[5], planes[5])
60 self.assertEqual(
61 eval(repr(schema), {"dtype": np.dtype, "MaskSchema": MaskSchema, "MaskPlane": MaskPlane}), schema
62 )
63 string = str(schema)
64 self.assertEqual(len(string.split("\n")), 17)
65 bit5 = schema.bit("M5")
66 self.assertIn(f"M5 [{bit5.index}@{hex(bit5.mask)}]: D5", string)
67 self.assertEqual(schema, MaskSchema(planes, np.uint8))
68 self.assertNotEqual(schema, MaskSchema(planes, np.int16))
69 self.assertNotEqual(schema, MaskSchema(planes[:-1], np.uint8))
70 self.assertEqual(schema.dtype, np.dtype(np.uint8))
71 self.assertEqual(schema.mask_size, 3)
72 self.assertEqual(schema.names, {f"M{i}" for i in range(17)})
73 self.assertEqual(schema.descriptions, {f"M{i}": f"D{i}" for i in range(17)})
74 bit7 = schema.bit("M7")
75 bitmask57 = schema.bitmask("M5", "M7")
76 self.assertTrue(bitmask57[bit5.index] & bit5.mask)
77 self.assertTrue(bitmask57[bit7.index] & bit7.mask)
78 bitmask57[bit5.index] &= ~bit5.mask
79 bitmask57[bit7.index] &= ~bit7.mask
80 self.assertFalse(bitmask57.any())
81 splits = schema.split(np.int16)
82 self.assertEqual(len(splits), 2)
83 self.assertEqual(splits[0].mask_size, 1)
84 self.assertEqual(splits[1].mask_size, 1)
85 self.assertEqual(list(splits[0]) + list(splits[1]), [p for p in planes if p is not None])
86 self.assertEqual(len(splits[0]), 15)
87 self.assertEqual(len(splits[1]), 2)
89 def test_schema_from_fits_header(self) -> None:
90 """MaskSchema.from_fits_header inverts update_header, assuming the
91 default uint8 dtype.
92 """
93 planes = [
94 MaskPlane("NO_DATA", "No data was available for this pixel."),
95 MaskPlane("COSMIC_RAY", "A cosmic ray affected this pixel."),
96 MaskPlane("DETECTED", "Pixel was part of a detected source."),
97 ]
98 schema = MaskSchema(planes, dtype=np.uint8)
99 header = astropy.io.fits.Header()
100 schema.update_header(header)
101 result = MaskSchema.from_fits_header(header)
102 self.assertEqual(result.dtype, np.dtype(np.uint8))
103 self.assertEqual(list(result), planes)
104 self.assertEqual(result, schema)
106 def test_schema_from_fits_header_preserves_gaps(self) -> None:
107 """A None placeholder between planes is reconstructed from the gap in
108 the MSKN card numbering.
109 """
110 planes: list[MaskPlane | None] = [MaskPlane("A", "a"), None, MaskPlane("B", "b")]
111 header = astropy.io.fits.Header()
112 MaskSchema(planes, dtype=np.uint8).update_header(header)
113 self.assertEqual(list(MaskSchema.from_fits_header(header)), planes)
115 def test_schema_from_fits_header_requires_cards(self) -> None:
116 """A header with no MSKN cards cannot describe a schema."""
117 with self.assertRaises(ValueError):
118 MaskSchema.from_fits_header(astropy.io.fits.Header())
120 def test_basics(self) -> None:
121 """Test some basic Mask functionality."""
122 planes = self.make_mask_planes(35, n_placeholders=5)
123 schema = MaskSchema(planes, dtype=np.uint8)
124 bbox = Box.factory[5:50, 6:60]
125 mask = Mask(
126 0,
127 schema=schema,
128 bbox=bbox,
129 metadata={"four_and_a_half": 4.5},
130 )
132 self.assertIs(mask[...], mask)
133 self.assertEqual(mask.__eq__(42), NotImplemented)
134 self.assertEqual(mask, mask)
135 self.maxDiff = None
136 self.assertEqual(
137 str(mask),
138 "Mask([y=5:50, x=6:60], ['M34', 'M15', 'M29', 'M1', 'M20', 'M11', 'M13', 'M7', 'M17', 'M12', "
139 "'M31', 'M16', 'M2', 'M3', 'M8', 'M26', 'M22', 'M5', 'M18', 'M19', 'M24', 'M21', 'M27', 'M6', "
140 "'M28', 'M10', 'M4', 'M23', 'M0', 'M25', 'M9', 'M14', 'M33', 'M32', 'M30'])",
141 )
142 self.assertTrue(
143 repr(mask).startswith(
144 "Mask(..., bbox=Box(y=Interval(start=5, stop=50), x=Interval(start=6, stop=60)), "
145 "schema=MaskSchema([MaskPlane(name='M34', description='D34')"
146 ),
147 f"Repr: {mask!r}",
148 )
150 with self.assertRaises(TypeError):
151 # No bbox, size or array.
152 Mask(0, schema=schema)
154 with self.assertRaises(ValueError):
155 # Box mismatch.
156 Mask(mask.array, schema=schema, bbox=Box.factory[0:20, -5:45])
158 with self.assertRaises(ValueError):
159 # Shape mismatch.
160 Mask(mask.array, schema=schema, shape=(5, 10, 5))
162 with self.assertRaises(ValueError):
163 # Cannot be 2-D.
164 Mask(mask.array.reshape((2430, 5)), schema=schema, bbox=Box.factory[0:20, -5:45])
166 def test_read_write(self) -> None:
167 """Explicit calls to read and write fits."""
168 planes = self.make_mask_planes(35, n_placeholders=5)
169 schema = MaskSchema(planes, dtype=np.uint8)
170 bbox = Box.factory[5:50, 6:60]
171 mask = Mask(
172 0,
173 schema=schema,
174 bbox=bbox,
175 metadata={"four_and_a_half": 4.5},
176 )
177 with lsst.utils.tests.getTempFilePath(".fits") as tmpFile:
178 mask.write(tmpFile)
180 new = Mask.read(tmpFile)
181 self.assertEqual(new, mask)
182 # __eq__ ignores metadata.
183 self.assertEqual(new.metadata["four_and_a_half"], 4.5)
184 self.assertEqual(new.metadata, mask.metadata)
186 def test_serialize_multi(self) -> None:
187 """Test serializing a mask with more than 31 mask planes, requiring
188 more than one HDU and EXTVER.
190 Note that serialization for simpler cases is covered by
191 test_masked_image.py.
192 """
193 planes = self.make_mask_planes(35, n_placeholders=5)
194 schema = MaskSchema(planes, dtype=np.uint8)
195 bbox = Box.factory[5:50, 6:60]
196 mask = Mask(0, schema=schema, bbox=bbox, metadata={"four_and_a_half": 4.5})
197 shape = bbox.shape
198 for plane in schema:
199 if plane is not None:
200 mask.set(plane.name, self.rng.random(shape) > 0.5)
201 with RoundtripFits(self, mask) as roundtrip:
202 fits = roundtrip.inspect()
203 self.assertEqual(fits[1].header["EXTNAME"], "MASK")
204 self.assertEqual(fits[1].header.get("EXTVER", 1), 1)
205 self.assertEqual(fits[1].header["ZCMPTYPE"], "GZIP_2")
206 self.assertEqual(fits[2].header["EXTNAME"], "MASK")
207 self.assertEqual(fits[2].header["EXTVER"], 2)
208 self.assertEqual(fits[2].header["ZCMPTYPE"], "GZIP_2")
209 n = 0
210 for plane in planes:
211 if plane is not None:
212 hdu = fits[1] if n < 31 else fits[2]
213 self.assertEqual(hdu.header[f"MSKN{(n % 31):04d}"], plane.name)
214 self.assertEqual(hdu.header[f"MSKM{(n % 31):04d}"], 1 << (n % 31))
215 self.assertEqual(hdu.header[f"MSKD{(n % 31):04d}"], plane.description)
216 n += 1
217 assert_masks_equal(self, mask, roundtrip.result)
219 def test_add_plane_returns_new_mask(self) -> None:
220 """Adding a plane returns a new mask, leaves the original (and any
221 views of it) untouched, and always reallocates the backing array.
222 """
223 planes = self.make_mask_planes(3, n_placeholders=0)
224 schema = MaskSchema(planes, dtype=np.uint8)
225 bbox = Box.factory[5:50, 6:60]
226 mask = Mask(0, schema=schema, bbox=bbox)
227 m0 = self.rng.random(bbox.shape) > 0.5
228 mask.set("M0", m0)
229 view = mask[bbox] # shares the array and old schema with mask
230 original_array = mask.array
232 new_mask = mask.add_plane("OUTSIDE_STENCIL", "Pixel lies outside the stencil.")
234 # The original mask and any views keep the old schema and array.
235 self.assertNotIn("OUTSIDE_STENCIL", mask.schema.names)
236 self.assertNotIn("OUTSIDE_STENCIL", view.schema.names)
237 self.assertIs(mask.array, original_array)
238 # The new mask reallocated a fresh array and carries the new plane.
239 self.assertIsNot(new_mask.array, original_array)
240 self.assertIn("OUTSIDE_STENCIL", new_mask.schema.names)
241 self.assertEqual(new_mask.schema.descriptions["OUTSIDE_STENCIL"], "Pixel lies outside the stencil.")
242 # The new plane is the fourth (overall index 3) so it lives in byte 0.
243 bit = new_mask.schema.bit("OUTSIDE_STENCIL")
244 self.assertEqual(bit.index, 0)
245 self.assertEqual(bit.mask, 1 << 3)
246 self.assertEqual(new_mask.schema.mask_size, 1)
247 # Existing plane data is preserved and the new plane starts all-False.
248 np.testing.assert_array_equal(new_mask.get("M0"), m0)
249 self.assertFalse(new_mask.get("OUTSIDE_STENCIL").any())
251 def test_add_plane_grows_byte(self) -> None:
252 """Adding a ninth plane (crossing the 8-plane boundary) gives the new
253 mask an extra byte while preserving existing plane data.
254 """
255 planes = self.make_mask_planes(8, n_placeholders=0)
256 schema = MaskSchema(planes, dtype=np.uint8)
257 bbox = Box.factory[5:50, 6:60]
258 mask = Mask(0, schema=schema, bbox=bbox)
259 set_planes = {}
260 for plane in planes:
261 assert plane is not None
262 boolean_mask = self.rng.random(bbox.shape) > 0.5
263 mask.set(plane.name, boolean_mask)
264 set_planes[plane.name] = boolean_mask
266 new_mask = mask.add_plane("OUTSIDE_STENCIL", "Pixel lies outside the stencil.")
268 # The original is unchanged; the new mask spills into a second byte.
269 self.assertEqual(mask.schema.mask_size, 1)
270 bit = new_mask.schema.bit("OUTSIDE_STENCIL")
271 self.assertEqual(bit.index, 1)
272 self.assertEqual(bit.mask, 1 << 0)
273 self.assertEqual(new_mask.schema.mask_size, 2)
274 self.assertEqual(new_mask.array.shape, bbox.shape + (2,))
275 self.assertFalse(new_mask.get("OUTSIDE_STENCIL").any())
276 # Every pre-existing plane keeps its data.
277 for name, boolean_mask in set_planes.items():
278 np.testing.assert_array_equal(new_mask.get(name), boolean_mask)
280 def test_add_planes_multiple(self) -> None:
281 """Several planes can be added in a single call."""
282 planes = self.make_mask_planes(3, n_placeholders=0)
283 bbox = Box.factory[0:4, 0:5]
284 mask = Mask(0, schema=MaskSchema(planes, dtype=np.uint8), bbox=bbox)
285 m0 = self.rng.random(bbox.shape) > 0.5
286 mask.set("M0", m0)
288 new_mask = mask.add_planes([MaskPlane("A", "plane a"), MaskPlane("B", "plane b")])
290 self.assertEqual(set(mask.schema.names), {"M0", "M1", "M2"}) # original unchanged
291 self.assertEqual(set(new_mask.schema.names), {"M0", "M1", "M2", "A", "B"})
292 np.testing.assert_array_equal(new_mask.get("M0"), m0)
293 self.assertFalse(new_mask.get("A").any())
294 self.assertFalse(new_mask.get("B").any())
296 def test_add_planes_drop_reassigns_bits(self) -> None:
297 """Dropping a plane compacts the schema, reassigns the planes after it
298 to lower bits, and repacks the retained pixel values by name.
299 """
300 bbox = Box.factory[0:4, 0:5]
301 schema = MaskSchema([MaskPlane("A", "a"), MaskPlane("B", "b"), MaskPlane("C", "c")], dtype=np.uint8)
302 mask = Mask(0, schema=schema, bbox=bbox)
303 a = self.rng.random(bbox.shape) > 0.5
304 c = self.rng.random(bbox.shape) > 0.5
305 mask.set("A", a)
306 mask.set("B", self.rng.random(bbox.shape) > 0.5)
307 mask.set("C", c)
309 new_mask = mask.add_planes([MaskPlane("D", "d")], drop=["B"])
311 # B is gone; D is appended after the retained planes.
312 self.assertEqual(list(new_mask.schema.names), ["A", "C", "D"])
313 self.assertNotIn("B", new_mask.schema.names)
314 # C moved down from bit 2 to bit 1; D takes bit 2.
315 self.assertEqual(new_mask.schema.bit("A").mask, 1 << 0)
316 self.assertEqual(new_mask.schema.bit("C").mask, 1 << 1)
317 self.assertEqual(new_mask.schema.bit("D").mask, 1 << 2)
318 # Retained pixel values follow their planes; the new plane is cleared.
319 np.testing.assert_array_equal(new_mask.get("A"), a)
320 np.testing.assert_array_equal(new_mask.get("C"), c)
321 self.assertFalse(new_mask.get("D").any())
323 def test_add_planes_with_placeholder(self) -> None:
324 """``None`` placeholders reserve bits. A pre-existing placeholder
325 keeps its position, and a ``None`` interleaved in the added planes
326 stays where it is placed rather than moving to the end; both survive
327 a round-trip.
328 """
329 bbox = Box.factory[0:4, 0:5]
330 # Schema with a pre-existing placeholder reserving bit 1.
331 schema = MaskSchema([MaskPlane("A", "a"), None, MaskPlane("B", "b")], dtype=np.uint8)
332 mask = Mask(0, schema=schema, bbox=bbox)
333 a = self.rng.random(bbox.shape) > 0.5
334 b = self.rng.random(bbox.shape) > 0.5
335 mask.set("A", a)
336 mask.set("B", b)
338 # Append a block that itself contains an interior placeholder.
339 new_mask = mask.add_planes([MaskPlane("C", "c"), None, MaskPlane("D", "d")])
341 # The pre-existing placeholder stays at bit 1; the added placeholder
342 # stays between C and D (bit 4), not at the end.
343 self.assertEqual(
344 list(new_mask.schema),
345 [MaskPlane("A", "a"), None, MaskPlane("B", "b"), MaskPlane("C", "c"), None, MaskPlane("D", "d")],
346 )
347 self.assertEqual(new_mask.schema.bit("A").mask, 1 << 0)
348 self.assertEqual(new_mask.schema.bit("B").mask, 1 << 2)
349 self.assertEqual(new_mask.schema.bit("C").mask, 1 << 3)
350 self.assertEqual(new_mask.schema.bit("D").mask, 1 << 5)
351 # Retained pixel values follow their planes; new planes start cleared.
352 np.testing.assert_array_equal(new_mask.get("A"), a)
353 np.testing.assert_array_equal(new_mask.get("B"), b)
354 self.assertFalse(new_mask.get("C").any())
355 self.assertFalse(new_mask.get("D").any())
357 with RoundtripFits(self, new_mask) as roundtrip:
358 assert_masks_equal(self, new_mask, roundtrip.result)
360 def test_add_planes_drop_unknown_raises(self) -> None:
361 """Dropping a plane that does not exist is an error."""
362 mask = Mask(0, schema=MaskSchema([MaskPlane("A", "a")], dtype=np.uint8), bbox=Box.factory[0:2, 0:2])
363 with self.assertRaises(ValueError):
364 mask.add_planes([], drop=["NOPE"])
366 def test_add_plane_duplicate_raises(self) -> None:
367 """Adding a plane whose name already exists is an error."""
368 planes = self.make_mask_planes(3, n_placeholders=0)
369 schema = MaskSchema(planes, dtype=np.uint8)
370 mask = Mask(0, schema=schema, bbox=Box.factory[0:4, 0:4])
371 with self.assertRaises(ValueError):
372 mask.add_plane("M0", "Duplicate of an existing plane.")
374 def test_add_plane_roundtrip(self) -> None:
375 """A runtime-added plane and its data survive a FITS round trip."""
376 planes = self.make_mask_planes(8, n_placeholders=0)
377 schema = MaskSchema(planes, dtype=np.uint8)
378 bbox = Box.factory[5:50, 6:60]
379 mask = Mask(0, schema=schema, bbox=bbox)
380 mask = mask.add_plane("OUTSIDE_STENCIL", "Pixel lies outside the stencil.")
381 mask.set("OUTSIDE_STENCIL", self.rng.random(bbox.shape) > 0.5)
382 with lsst.utils.tests.getTempFilePath(".fits") as tmpFile:
383 mask.write(tmpFile)
384 new = Mask.read(tmpFile)
385 self.assertEqual(new, mask)
386 self.assertEqual(new.schema.descriptions["OUTSIDE_STENCIL"], "Pixel lies outside the stencil.")
387 assert_masks_equal(self, new, mask)
389 def test_legacy_non_cell_coadd_plane_map(self) -> None:
390 """The non-cell coadd map defines a distinct ``SENSOR_EDGE`` plane."""
391 plane_map = get_legacy_non_cell_coadd_mask_planes()
392 self.assertIn("SENSOR_EDGE", plane_map)
393 self.assertEqual(plane_map["SENSOR_EDGE"].name, "SENSOR_EDGE")
395 def test_guess_legacy_plane_map_coadd_discriminator(self) -> None:
396 """``INEXACT_PSF`` routes to a coadd map; ``SENSOR_EDGE`` distinguishes
397 non-cell coadds (which use it) from cell coadds (which use CELL_EDGE).
398 """
399 non_cell = _guess_legacy_plane_map({"INEXACT_PSF": 11, "SENSOR_EDGE": 14})
400 self.assertIn("SENSOR_EDGE", non_cell)
401 cell = _guess_legacy_plane_map({"INEXACT_PSF": 11})
402 self.assertNotIn("SENSOR_EDGE", cell)
404 @unittest.skipUnless(DATA_DIR is not None, "TESTDATA_IMAGES_DIR is not in the environment.")
405 def test_legacy(self) -> None:
406 """Test Mask.read_legacy, Mask.to_legacy, and Mask.from_legacy."""
407 assert DATA_DIR is not None, "Guaranteed by decorator."
408 filename = os.path.join(DATA_DIR, "dp2", "legacy", "visit_image.fits")
409 plane_map = get_legacy_visit_image_mask_planes()
410 mask = Mask.read_legacy(filename, ext=2, plane_map=plane_map)
411 self.assertEqual(mask.schema.names, {p.name for p in plane_map.values()})
412 try:
413 from lsst.afw.image import MaskedImageFitsReader
414 except ImportError:
415 raise unittest.SkipTest("'lsst.afw.image' could not be imported.") from None
416 reader = MaskedImageFitsReader(filename)
417 self.assertEqual(mask.bbox, Box.from_legacy(reader.readBBox()))
418 legacy_mask = reader.readMask()
419 compare_mask_to_legacy(self, mask, legacy_mask, plane_map)
420 compare_mask_to_legacy(self, mask, mask.to_legacy(plane_map), plane_map)
421 assert_masks_equal(self, mask, Mask.from_legacy(legacy_mask, plane_map=plane_map))
422 # Write the mask out in the new format, and test that we can read it
423 # back either way.
424 with RoundtripFits(self, mask, storage_class="MaskV2") as roundtrip:
425 with self.subTest():
426 try:
427 import lsst.afw.image
428 except ImportError:
429 raise unittest.SkipTest("afw could not be imported") from None
430 legacy_mask = roundtrip.get(storageClass="Mask")
431 self.assertIsInstance(legacy_mask, lsst.afw.image.Mask)
432 compare_mask_to_legacy(self, mask, legacy_mask)
433 assert_masks_equal(self, roundtrip.result, mask)
436if __name__ == "__main__":
437 unittest.main()