Coverage for tests / test_masked_image.py: 24%

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 

21 

22from lsst.images import Box, Image, MaskedImage, MaskPlane, MaskSchema, get_legacy_visit_image_mask_planes 

23from lsst.images.fits import FitsCompressionOptions 

24from lsst.images.tests import RoundtripFits, assert_masked_images_equal, compare_masked_image_to_legacy 

25 

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

27 

28 

29class MaskedImageTestCase(unittest.TestCase): 

30 """Tests for the MaskedImage class and the basics of the archive system.""" 

31 

32 def setUp(self) -> None: 

33 self.maxDiff = None 

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

35 self.masked_image = MaskedImage( 

36 Image(self.rng.normal(100.0, 8.0, size=(200, 251)), dtype=np.float64, unit=u.nJy, start=(5, 8)), 

37 mask_schema=MaskSchema( 

38 [ 

39 MaskPlane("BAD", "Pixel is very bad, possibly downright evil."), 

40 MaskPlane("HUNGRY", "Pixel hasn't had enough to eat today."), 

41 ] 

42 ), 

43 metadata={"fifty": "5 * 10"}, 

44 ) 

45 self.masked_image.mask.array |= np.multiply.outer( 

46 self.masked_image.image.array < 102.0, 

47 self.masked_image.mask.schema.bitmask("BAD"), 

48 ) 

49 self.masked_image.mask.array |= np.multiply.outer( 

50 self.masked_image.image.array > 98.0, 

51 self.masked_image.mask.schema.bitmask("HUNGRY"), 

52 ) 

53 self.masked_image.variance.array = self.rng.normal(64.0, 0.5, size=self.masked_image.bbox.shape) 

54 

55 def test_construction(self) -> None: 

56 """Test that the MaskedImage construction (in setUp) worked.""" 

57 self.assertEqual(self.masked_image.bbox, Box.factory[5:205, 8:259]) 

58 self.assertEqual(self.masked_image.mask.bbox, self.masked_image.bbox) 

59 self.assertEqual(self.masked_image.variance.bbox, self.masked_image.bbox) 

60 self.assertEqual(self.masked_image.image.array.shape, self.masked_image.bbox.shape) 

61 self.assertEqual(self.masked_image.mask.array.shape, self.masked_image.bbox.shape + (1,)) 

62 self.assertEqual(self.masked_image.variance.array.shape, self.masked_image.bbox.shape) 

63 self.assertEqual(self.masked_image.unit, u.nJy) 

64 self.assertEqual(self.masked_image.variance.unit, u.nJy**2) 

65 self.assertEqual(self.masked_image.metadata, {"fifty": "5 * 10"}) 

66 # The checks below are subject to the vagaries of the RNG, but we want 

67 # the seed to be such that they all pass, or other tests will be 

68 # weaker. 

69 self.assertGreater( 

70 np.sum(self.masked_image.mask.array == self.masked_image.mask.schema.bitmask("BAD")), 0 

71 ) 

72 self.assertGreater( 

73 np.sum(self.masked_image.mask.array == self.masked_image.mask.schema.bitmask("HUNGRY")), 0 

74 ) 

75 self.assertGreater( 

76 np.sum(self.masked_image.mask.array == self.masked_image.mask.schema.bitmask("BAD", "HUNGRY")), 0 

77 ) 

78 

79 def test_fits_roundtrip(self) -> None: 

80 """Test that we can round-trip the MaskedImage through FITS, including 

81 subimage reads. 

82 """ 

83 subbox = Box.factory[11:20, 25:30] 

84 subslices = (slice(6, 15), slice(17, 22)) 

85 np.testing.assert_array_equal( 

86 self.masked_image.image.array[subslices], self.masked_image.image[subbox].array 

87 ) 

88 with RoundtripFits(self, self.masked_image, "MaskedImageV2") as roundtrip: 

89 subimage = roundtrip.get(bbox=subbox) 

90 # Check that we used lossless compression (the default). 

91 fits = roundtrip.inspect() 

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

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

94 self.assertEqual(fits[3].header["ZCMPTYPE"], "GZIP_2") 

95 assert_masked_images_equal(self, roundtrip.result, self.masked_image, expect_view=False) 

96 assert_masked_images_equal(self, subimage, roundtrip.result[subbox], expect_view=False) 

97 

98 def test_fits_roundtrip_lossy(self) -> None: 

99 """Test that we can round-trip the MaskedImage through FITS, including 

100 subimage reads, with lossy compression. 

101 """ 

102 subbox = Box.factory[11:20, 25:30] 

103 subslices = (slice(6, 15), slice(17, 22)) 

104 np.testing.assert_array_equal( 

105 self.masked_image.image.array[subslices], self.masked_image.image[subbox].array 

106 ) 

107 with tempfile.NamedTemporaryFile(suffix=".fits", delete_on_close=False, delete=True) as tmp: 

108 tmp.close() 

109 self.masked_image.write_fits( 

110 tmp.name, 

111 image_compression=FitsCompressionOptions.LOSSY, 

112 variance_compression=FitsCompressionOptions.LOSSY, 

113 compression_seed=50, 

114 ) 

115 roundtripped = MaskedImage.read_fits(tmp.name) 

116 subimage = MaskedImage.read_fits(tmp.name, bbox=subbox) 

117 with astropy.io.fits.open(tmp.name, disable_image_compression=True) as fits: 

118 self.assertEqual(fits[1].header["ZCMPTYPE"], "RICE_1") 

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

120 self.assertEqual(fits[3].header["ZCMPTYPE"], "RICE_1") 

121 assert_masked_images_equal(self, roundtripped, self.masked_image, expect_view=False, rtol=0.01) 

122 assert_masked_images_equal(self, subimage, roundtripped[subbox], expect_view=False) 

123 

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

125 def test_legacy(self) -> None: 

126 """Test MaskedImage.read_legacy, MaskedImage.to_legacy, and 

127 MaskedImage.from_legacy. 

128 """ 

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

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

131 plane_map = get_legacy_visit_image_mask_planes() 

132 masked_image = MaskedImage.read_legacy(filename, plane_map=plane_map) 

133 try: 

134 from lsst.afw.image import MaskedImageFitsReader 

135 except ImportError: 

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

137 reader = MaskedImageFitsReader(filename) 

138 legacy_masked_image = reader.read() 

139 compare_masked_image_to_legacy( 

140 self, masked_image, legacy_masked_image, plane_map=plane_map, expect_view=False 

141 ) 

142 compare_masked_image_to_legacy( 

143 self, 

144 masked_image, 

145 masked_image.to_legacy(plane_map=plane_map), 

146 plane_map=plane_map, 

147 expect_view=True, 

148 ) 

149 compare_masked_image_to_legacy( 

150 self, 

151 MaskedImage.from_legacy(legacy_masked_image, plane_map=plane_map), 

152 legacy_masked_image, 

153 expect_view=True, 

154 plane_map=plane_map, 

155 ) 

156 

157 

158if __name__ == "__main__": 

159 unittest.main()