Coverage for tests/test_masked_image.py: 89%

158 statements  

« prev     ^ index     » next       coverage.py v7.14.3, created at 2026-07-09 02:27 -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. 

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 ( 

25 RoundtripFits, 

26 RoundtripJson, 

27 RoundtripNdf, 

28 assert_masked_images_equal, 

29 compare_masked_image_to_legacy, 

30) 

31 

32try: 

33 import h5py # noqa: F401 

34 

35 HAVE_H5PY = True 

36except ImportError: 

37 HAVE_H5PY = False 

38 

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

40 

41 

42class MaskedImageTestCase(unittest.TestCase): 

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

44 

45 def setUp(self) -> None: 

46 self.maxDiff = None 

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

48 self.masked_image = MaskedImage( 

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

50 mask_schema=MaskSchema( 

51 [ 

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

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

54 ] 

55 ), 

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

57 ) 

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

59 self.masked_image.image.array < 102.0, 

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

61 ) 

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

63 self.masked_image.image.array > 98.0, 

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

65 ) 

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

67 

68 def test_construction(self) -> None: 

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

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

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

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

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

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

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

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

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

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

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

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

81 # weaker. 

82 self.assertGreater( 

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

84 ) 

85 self.assertGreater( 

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

87 ) 

88 self.assertGreater( 

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

90 ) 

91 

92 self.assertIs(self.masked_image[...], self.masked_image) 

93 self.assertEqual( 

94 str(self.masked_image), "MaskedImage(Image([y=5:205, x=8:259], float64), ['BAD', 'HUNGRY'])" 

95 ) 

96 self.assertEqual( 

97 repr(self.masked_image), 

98 "MaskedImage(Image(..., bbox=Box(y=Interval(start=5, stop=205), x=Interval(start=8, stop=259)), " 

99 "dtype=dtype('float64')), mask_schema=MaskSchema([MaskPlane(name='BAD', description='Pixel is " 

100 "very bad, possibly downright evil.'), MaskPlane(name='HUNGRY', description=\"Pixel hasn't had " 

101 "enough to eat today.\")], dtype=dtype('uint8')))", 

102 ) 

103 copy = self.masked_image.copy() 

104 original = self.masked_image.image.array[0, 0] 

105 copy.image.array[0, 0] = 38.0 

106 self.assertEqual(self.masked_image.image.array[0, 0], original) 

107 self.assertEqual(copy.image.array[0, 0], 38.0) 

108 

109 # Test error conditions. 

110 with self.assertRaises(ValueError): 

111 # Disagreement over mask bbox. 

112 MaskedImage(Image(42.0, shape=(5, 6)), mask=self.masked_image.mask) 

113 with self.assertRaises(TypeError): 

114 # No mask definition. 

115 MaskedImage(self.masked_image.image, variance=self.masked_image.variance) 

116 with self.assertRaises(TypeError): 

117 # Can not provide mask and mask schema. 

118 MaskedImage( 

119 Image(42.0, shape=(5, 5)), 

120 mask=self.masked_image.mask, 

121 mask_schema=self.masked_image.mask.schema, 

122 ) 

123 with self.assertRaises(ValueError): 

124 # image and variance bbox disagreement. 

125 MaskedImage( 

126 Image(42.0, shape=(5, 5)), 

127 mask_schema=self.masked_image.mask.schema, 

128 variance=self.masked_image.variance, 

129 ) 

130 with self.assertRaises(ValueError): 

131 # no image unit but there is variance unit. 

132 MaskedImage( 

133 Image(42.0, shape=(5, 5)), 

134 mask_schema=self.masked_image.mask.schema, 

135 variance=Image(1.0, shape=(5, 5), unit=u.nJy), 

136 ) 

137 with self.assertRaises(ValueError): 

138 # image and variance units disagree. 

139 MaskedImage( 

140 Image(42.0, shape=(5, 5), unit=u.nJy), 

141 mask_schema=self.masked_image.mask.schema, 

142 variance=Image(1.0, shape=(5, 5), unit=u.nJy), 

143 ) 

144 

145 def test_subset(self) -> None: 

146 """Test assignment of subset.""" 

147 copy = self.masked_image.copy() 

148 subset = copy.local[0:10, 20:30].copy() 

149 subset.image[...] = Image(42.0, shape=(10, 10), unit=u.nJy) 

150 copy[subset.bbox] = subset 

151 self.assertEqual(copy.image.array[0, 20], 42.0) 

152 self.assertEqual(copy.image.array[0, 0], self.masked_image.image.array[0, 0]) 

153 

154 def test_mask_setter(self) -> None: 

155 """The mask plane can be replaced, e.g. with one grown by add_plane.""" 

156 masked_image = self.masked_image.copy() 

157 bad = masked_image.mask.get("BAD") 

158 masked_image.mask = masked_image.mask.add_plane("OUTSIDE_STENCIL", "Pixel lies outside the stencil.") 

159 self.assertIn("OUTSIDE_STENCIL", masked_image.mask.schema.names) 

160 self.assertEqual(masked_image.mask.bbox, masked_image.image.bbox) 

161 np.testing.assert_array_equal(masked_image.mask.get("BAD"), bad) 

162 self.assertFalse(masked_image.mask.get("OUTSIDE_STENCIL").any()) 

163 # A mask whose bounding box disagrees with the image is rejected. 

164 with self.assertRaises(ValueError): 

165 masked_image.mask = masked_image.mask[Box.factory[10:20, 12:22]] 

166 

167 def test_fits_roundtrip(self) -> None: 

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

169 subimage reads. 

170 """ 

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

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

173 np.testing.assert_array_equal( 

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

175 ) 

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

177 subimage = roundtrip.get(bbox=subbox) 

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

179 fits = roundtrip.inspect() 

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

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

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

183 # Test reading back in as a legacy MaskedImage. 

184 with self.subTest(): 

185 try: 

186 import lsst.afw.image 

187 except ImportError: 

188 raise unittest.SkipTest("afw could not be imported") from None 

189 legacy_masked_image = roundtrip.get(storageClass="MaskedImage") 

190 self.assertIsInstance(legacy_masked_image, lsst.afw.image.MaskedImage) 

191 compare_masked_image_to_legacy( 

192 self, self.masked_image, legacy_masked_image, expect_view=False 

193 ) 

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

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

196 

197 def test_fits_roundtrip_lossy(self) -> None: 

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

199 subimage reads, with lossy compression. 

200 """ 

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

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

203 np.testing.assert_array_equal( 

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

205 ) 

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

207 tmp.close() 

208 self.masked_image.write( 

209 tmp.name, 

210 compression_options={ 

211 "image": FitsCompressionOptions.LOSSY, 

212 "variance": FitsCompressionOptions.LOSSY, 

213 }, 

214 compression_seed=50, 

215 ) 

216 roundtripped = MaskedImage.read(tmp.name) 

217 subimage = MaskedImage.read(tmp.name, bbox=subbox) 

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

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

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

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

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

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

224 

225 @unittest.skipUnless(HAVE_H5PY, "h5py is not installed") 

226 def test_round_trip_ndf_compatible_mask(self): 

227 """NDF round-trip for the default-setup MaskedImage (2 planes ≤ 8).""" 

228 with RoundtripNdf(self, self.masked_image, "MaskedImageV2") as roundtrip: 

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

230 

231 @unittest.skipUnless(HAVE_H5PY, "h5py is not installed") 

232 def test_round_trip_ndf_incompatible_mask(self): 

233 """NDF round-trip for a >8-plane mask (uses native 3D mask array, 

234 to MORE/LSST/MASK). 

235 """ 

236 rng = np.random.default_rng(7) 

237 planes = [MaskPlane(f"P{i}", f"plane {i}") for i in range(12)] 

238 wide = MaskedImage( 

239 Image( 

240 rng.normal(100.0, 8.0, size=(50, 60)), 

241 dtype=np.float64, 

242 unit=u.nJy, 

243 yx0=(0, 0), 

244 ), 

245 mask_schema=MaskSchema(planes), 

246 ) 

247 wide.variance.array = rng.normal(64.0, 0.5, size=wide.bbox.shape) 

248 with RoundtripNdf(self, wide, "MaskedImageV2") as roundtrip: 

249 assert_masked_images_equal(self, roundtrip.result, wide, expect_view=False) 

250 

251 @unittest.skipUnless(HAVE_H5PY, "h5py is not installed") 

252 def test_round_trip_ndf_many_plane_mask(self): 

253 """NDF round-trip for a mask that needs more than one int32 chunk.""" 

254 rng = np.random.default_rng(11) 

255 planes = [MaskPlane(f"P{i}", f"plane {i}") for i in range(40)] 

256 wide = MaskedImage( 

257 Image( 

258 rng.normal(100.0, 8.0, size=(10, 12)), 

259 dtype=np.float64, 

260 unit=u.nJy, 

261 yx0=(0, 0), 

262 ), 

263 mask_schema=MaskSchema(planes), 

264 ) 

265 wide.mask.set("P0", wide.image.array > 100.0) 

266 wide.mask.set("P17", wide.image.array < 95.0) 

267 wide.mask.set("P39", wide.image.array > 110.0) 

268 wide.variance.array = rng.normal(64.0, 0.5, size=wide.bbox.shape) 

269 with RoundtripNdf(self, wide, "MaskedImageV2") as roundtrip: 

270 assert_masked_images_equal(self, roundtrip.result, wide, expect_view=False) 

271 

272 @unittest.skipUnless(HAVE_H5PY, "h5py is not installed") 

273 def test_fits_ndf_consistency(self): 

274 """FITS and NDF backends produce equal MaskedImages on round-trip.""" 

275 with ( 

276 RoundtripFits(self, self.masked_image) as fits_rt, 

277 RoundtripNdf(self, self.masked_image) as ndf_rt, 

278 ): 

279 assert_masked_images_equal(self, self.masked_image, fits_rt.result, expect_view=False) 

280 assert_masked_images_equal(self, self.masked_image, ndf_rt.result, expect_view=False) 

281 assert_masked_images_equal(self, fits_rt.result, ndf_rt.result, expect_view=False) 

282 

283 def test_fits_json_consistency(self): 

284 """FITS and JSON backends produce equal MaskedImages on round-trip.""" 

285 with ( 

286 RoundtripFits(self, self.masked_image) as fits_rt, 

287 RoundtripJson(self, self.masked_image) as json_rt, 

288 ): 

289 assert_masked_images_equal(self, self.masked_image, fits_rt.result, expect_view=False) 

290 assert_masked_images_equal(self, self.masked_image, json_rt.result, expect_view=False) 

291 assert_masked_images_equal(self, fits_rt.result, json_rt.result, expect_view=False) 

292 

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

294 def test_legacy(self) -> None: 

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

296 MaskedImage.from_legacy. 

297 """ 

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

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

300 plane_map = get_legacy_visit_image_mask_planes() 

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

302 try: 

303 from lsst.afw.image import MaskedImageFitsReader 

304 except ImportError: 

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

306 reader = MaskedImageFitsReader(filename) 

307 legacy_masked_image = reader.read() 

308 compare_masked_image_to_legacy( 

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

310 ) 

311 compare_masked_image_to_legacy( 

312 self, 

313 masked_image, 

314 masked_image.to_legacy(plane_map=plane_map), 

315 plane_map=plane_map, 

316 expect_view=True, 

317 ) 

318 compare_masked_image_to_legacy( 

319 self, 

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

321 legacy_masked_image, 

322 expect_view=True, 

323 plane_map=plane_map, 

324 ) 

325 

326 

327if __name__ == "__main__": 

328 unittest.main()