Coverage for tests / test_mask.py: 13%

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 unittest 

16 

17import numpy as np 

18 

19from lsst.images import Box, Mask, MaskPlane, MaskSchema, get_legacy_visit_image_mask_planes 

20from lsst.images.tests import RoundtripFits, assert_masks_equal, compare_mask_to_legacy 

21 

22DATA_DIR = os.environ.get("TESTDATA_IMAGES_DIR", None) 

23 

24 

25class MaskTestCase(unittest.TestCase): 

26 """Tests for Mask and its helper classes.""" 

27 

28 def setUp(self) -> None: 

29 self.maxDiff = None 

30 self.rng = np.random.default_rng(500) 

31 

32 def make_mask_planes(self, n_planes: int, n_placeholders: int) -> list[MaskPlane | None]: 

33 planes: list[MaskPlane | None] = [] 

34 for i in range(n_planes): 

35 planes.append(MaskPlane(f"M{i}", f"D{i}")) 

36 planes.extend([None] * n_placeholders) 

37 self.rng.shuffle(planes) 

38 return planes 

39 

40 def test_schema(self) -> None: 

41 """Test MaskSchema.""" 

42 self.assertEqual(MaskSchema.bits_per_element(np.uint8), 8) 

43 planes = self.make_mask_planes(17, 5) 

44 with self.assertRaises(TypeError): 

45 MaskSchema.bits_per_element(np.float32) 

46 schema = MaskSchema(planes, dtype=np.uint8) 

47 self.assertEqual(list(schema), planes) 

48 self.assertEqual(len(schema), len(planes)) 

49 self.assertEqual(schema[5], planes[5]) 

50 self.assertEqual( 

51 eval(repr(schema), {"dtype": np.dtype, "MaskSchema": MaskSchema, "MaskPlane": MaskPlane}), schema 

52 ) 

53 string = str(schema) 

54 self.assertEqual(len(string.split("\n")), 17) 

55 bit5 = schema.bit("M5") 

56 self.assertIn(f"M5 [{bit5.index}@{hex(bit5.mask)}]: D5", string) 

57 self.assertEqual(schema, MaskSchema(planes, np.uint8)) 

58 self.assertNotEqual(schema, MaskSchema(planes, np.int16)) 

59 self.assertNotEqual(schema, MaskSchema(planes[:-1], np.uint8)) 

60 self.assertEqual(schema.dtype, np.dtype(np.uint8)) 

61 self.assertEqual(schema.mask_size, 3) 

62 self.assertEqual(schema.names, {f"M{i}" for i in range(17)}) 

63 self.assertEqual(schema.descriptions, {f"M{i}": f"D{i}" for i in range(17)}) 

64 bit7 = schema.bit("M7") 

65 bitmask57 = schema.bitmask("M5", "M7") 

66 self.assertTrue(bitmask57[bit5.index] & bit5.mask) 

67 self.assertTrue(bitmask57[bit7.index] & bit7.mask) 

68 bitmask57[bit5.index] &= ~bit5.mask 

69 bitmask57[bit7.index] &= ~bit7.mask 

70 self.assertFalse(bitmask57.any()) 

71 splits = schema.split(np.int16) 

72 self.assertEqual(len(splits), 2) 

73 self.assertEqual(splits[0].mask_size, 1) 

74 self.assertEqual(splits[1].mask_size, 1) 

75 self.assertEqual(list(splits[0]) + list(splits[1]), [p for p in planes if p is not None]) 

76 self.assertEqual(len(splits[0]), 15) 

77 self.assertEqual(len(splits[1]), 2) 

78 

79 def test_serialize_multi(self) -> None: 

80 """Test serializing a mask with more than 31 mask planes, requiring 

81 more than one HDU and EXTVER. 

82 

83 Note that serialization for simpler cases is covered by 

84 test_masked_image.py. 

85 """ 

86 planes = self.make_mask_planes(35, n_placeholders=5) 

87 schema = MaskSchema(planes, dtype=np.uint8) 

88 bbox = Box.factory[5:50, 6:60] 

89 mask = Mask(0, schema=schema, bbox=bbox) 

90 shape = bbox.shape 

91 for plane in schema: 

92 if plane is not None: 

93 mask.set(plane.name, self.rng.random(shape) > 0.5) 

94 with RoundtripFits(self, mask) as roundtrip: 

95 fits = roundtrip.inspect() 

96 self.assertEqual(fits[1].header["EXTNAME"], "MASK") 

97 self.assertEqual(fits[1].header.get("EXTVER", 1), 1) 

98 self.assertEqual(fits[1].header["ZCMPTYPE"], "GZIP_2") 

99 self.assertEqual(fits[2].header["EXTNAME"], "MASK") 

100 self.assertEqual(fits[2].header["EXTVER"], 2) 

101 self.assertEqual(fits[2].header["ZCMPTYPE"], "GZIP_2") 

102 n = 0 

103 for plane in planes: 

104 if plane is not None: 

105 hdu = fits[1] if n < 31 else fits[2] 

106 self.assertEqual(hdu.header[f"MSKN{(n % 31) + 1:04d}"], plane.name) 

107 self.assertEqual(hdu.header[f"MSKM{(n % 31) + 1:04d}"], 1 << (n % 31)) 

108 self.assertEqual(hdu.header[f"MSKD{(n % 31) + 1:04d}"], plane.description) 

109 n += 1 

110 assert_masks_equal(self, mask, roundtrip.result) 

111 

112 @unittest.skipUnless(DATA_DIR is not None, "TESTDATA_IMAGES_DIR is not in the environment.") 

113 def test_legacy(self) -> None: 

114 """Test Mask.read_legacy, Mask.to_legacy, and Mask.from_legacy.""" 

115 assert DATA_DIR is not None, "Guaranteed by decorator." 

116 filename = os.path.join(DATA_DIR, "dp2", "legacy", "visit_image.fits") 

117 plane_map = get_legacy_visit_image_mask_planes() 

118 mask = Mask.read_legacy(filename, ext=2, plane_map=plane_map) 

119 try: 

120 from lsst.afw.image import MaskedImageFitsReader 

121 except ImportError: 

122 raise unittest.SkipTest("'lsst.afw.image' could not be imported.") from None 

123 reader = MaskedImageFitsReader(filename) 

124 self.assertEqual(mask.bbox, Box.from_legacy(reader.readBBox())) 

125 legacy_mask = reader.readMask() 

126 compare_mask_to_legacy(self, mask, legacy_mask, plane_map) 

127 compare_mask_to_legacy(self, mask, mask.to_legacy(plane_map), plane_map) 

128 assert_masks_equal(self, mask, Mask.from_legacy(legacy_mask, plane_map=plane_map)) 

129 

130 

131if __name__ == "__main__": 

132 unittest.main()